Fix ML-KEM arithmetic timing paths
This commit is contained in:
@@ -1,8 +1,9 @@
|
||||
// comp_decomp_sync.v - ML-KEM coefficient compression/decompression
|
||||
//
|
||||
// Streaming: one coefficient per cycle through pipeline_reg.
|
||||
// Streaming valid/ready interface. Compression is intentionally iterative to
|
||||
// avoid inferring a long combinational divider for /Q.
|
||||
// mode=0: compress — round((2^d * x) / Q) mod 2^d
|
||||
// mode=1: decompress — round((Q * x) / 2^d) mod Q
|
||||
// mode=1: decompress — round((Q * x) / 2^d)
|
||||
// Uses round-half-up (round(2.5)=3, round(3.5)=4).
|
||||
// Integer arithmetic:
|
||||
// compress: (x * 2^d + Q/2) / Q, lower d bits as result
|
||||
@@ -24,52 +25,81 @@ module comp_decomp_sync (
|
||||
input ready_i
|
||||
);
|
||||
|
||||
// 2^d for d in {4,5,10,11}; fits in 12 bits (max 2048)
|
||||
wire [11:0] two_pow_d;
|
||||
assign two_pow_d = 12'd1 << d;
|
||||
localparam [11:0] Q_VAL = 12'(`Q);
|
||||
localparam [4:0] COMP_MSB = 5'd22; // max ((Q-1)<<11)+Q/2 is 23 bits
|
||||
|
||||
// Product: 12-bit * 12-bit = max 2048*3328 = 6,815,744 → 23 bits (use 24)
|
||||
wire [23:0] product;
|
||||
reg busy_r;
|
||||
reg [4:0] bit_idx_r;
|
||||
reg [4:0] d_r;
|
||||
reg [23:0] dividend_r;
|
||||
reg [23:0] quotient_r;
|
||||
reg [12:0] rem_r;
|
||||
reg [11:0] data_r;
|
||||
reg valid_r;
|
||||
|
||||
// compress: x * 2^d; decompress: x * Q
|
||||
assign product = mode ? {12'b0, coeff_in} * {12'b0, 12'(`Q)}
|
||||
: {12'b0, coeff_in} * {12'b0, two_pow_d};
|
||||
wire fire_i = valid_i && ready_o;
|
||||
|
||||
// Rounding offset: compress→Q/2=1664; decompress→2^(d-1)
|
||||
wire [11:0] round_off;
|
||||
assign round_off = mode ? (two_pow_d >> 1) : 12'd1664;
|
||||
// Accept a new command only when the iterative divider and output slot are
|
||||
// both free. Current users tie ready_i high and issue one coefficient at a
|
||||
// time, but this keeps the module from accepting work it cannot retain.
|
||||
assign ready_o = !busy_r && (!valid_r || ready_i);
|
||||
assign coeff_out = data_r;
|
||||
assign valid_o = valid_r;
|
||||
|
||||
// Dividend = product + round_off (max ~ 6,817,408 fits in 24 bits)
|
||||
wire [23:0] dividend;
|
||||
assign dividend = product + {12'b0, round_off};
|
||||
wire [23:0] comp_dividend = ({12'd0, coeff_in} << d) + 24'd1664;
|
||||
wire [23:0] decomp_product =
|
||||
{12'd0, coeff_in} +
|
||||
({12'd0, coeff_in} << 8) +
|
||||
({12'd0, coeff_in} << 10) +
|
||||
({12'd0, coeff_in} << 11);
|
||||
wire [23:0] decomp_rounded = decomp_product + (24'd1 << (d - 1'b1));
|
||||
wire [11:0] decomp_result = decomp_rounded >> d;
|
||||
|
||||
// Divisor: compress→Q=3329; decompress→2^d
|
||||
wire [11:0] divisor;
|
||||
assign divisor = mode ? two_pow_d : 12'(`Q);
|
||||
wire [12:0] rem_shift = {rem_r[11:0], dividend_r[bit_idx_r]};
|
||||
wire rem_ge_q = rem_shift >= {1'b0, Q_VAL};
|
||||
wire [12:0] rem_sub_q = rem_shift - {1'b0, Q_VAL};
|
||||
wire [12:0] rem_next = rem_ge_q ? rem_sub_q : rem_shift;
|
||||
wire [23:0] quot_next = quotient_r | (rem_ge_q ? (24'd1 << bit_idx_r) : 24'd0);
|
||||
wire [11:0] comp_mask = (12'd1 << d_r) - 12'd1;
|
||||
|
||||
// Integer division (both ops positive → floor)
|
||||
// Quotient is computed as 24b (widest operand) but fits in 12b
|
||||
/* verilator lint_off WIDTHTRUNC */
|
||||
wire [11:0] div_result;
|
||||
assign div_result = dividend / {12'b0, divisor};
|
||||
/* verilator lint_on WIDTHTRUNC */
|
||||
always @(posedge clk or negedge rst_n) begin
|
||||
if (!rst_n) begin
|
||||
busy_r <= 1'b0;
|
||||
bit_idx_r <= 5'd0;
|
||||
d_r <= 5'd0;
|
||||
dividend_r <= 24'd0;
|
||||
quotient_r <= 24'd0;
|
||||
rem_r <= 13'd0;
|
||||
data_r <= 12'd0;
|
||||
valid_r <= 1'b0;
|
||||
end else begin
|
||||
if (valid_r && ready_i)
|
||||
valid_r <= 1'b0;
|
||||
|
||||
// Final modular reduction
|
||||
// compress: lower d bits (mask with 2^d - 1)
|
||||
// decompress: result < Q always, but apply mod Q for safety
|
||||
wire [11:0] mod_result;
|
||||
assign mod_result = mode ? (div_result % 12'(`Q)) : (div_result & (two_pow_d - 1'b1));
|
||||
|
||||
// Pipeline through valid/ready stage
|
||||
pipeline_reg #(.DW(12)) u_pipe (
|
||||
.clk (clk),
|
||||
.rst_n (rst_n),
|
||||
.data_i (mod_result),
|
||||
.valid_i(valid_i),
|
||||
.ready_o(ready_o),
|
||||
.data_o (coeff_out),
|
||||
.valid_o(valid_o),
|
||||
.ready_i(ready_i)
|
||||
);
|
||||
if (busy_r) begin
|
||||
rem_r <= rem_next;
|
||||
quotient_r <= quot_next;
|
||||
if (bit_idx_r == 5'd0) begin
|
||||
busy_r <= 1'b0;
|
||||
data_r <= quot_next[11:0] & comp_mask;
|
||||
valid_r <= 1'b1;
|
||||
end else begin
|
||||
bit_idx_r <= bit_idx_r - 5'd1;
|
||||
end
|
||||
end else if (fire_i) begin
|
||||
if (mode) begin
|
||||
data_r <= decomp_result;
|
||||
valid_r <= 1'b1;
|
||||
end else begin
|
||||
busy_r <= 1'b1;
|
||||
bit_idx_r <= COMP_MSB;
|
||||
d_r <= d;
|
||||
dividend_r <= comp_dividend;
|
||||
quotient_r <= 24'd0;
|
||||
rem_r <= 13'd0;
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
endmodule
|
||||
|
||||
@@ -38,8 +38,9 @@ xvlog -sv ${NTT_DIR}/barrett_mul.v
|
||||
# ================================================================
|
||||
puts "=== Compiling PolyMul RTL sources ==="
|
||||
|
||||
# Basecase multiplier (instantiates barrett_mul)
|
||||
# Basecase multipliers
|
||||
xvlog -sv ${PM_DIR}/basecase_mul.v
|
||||
xvlog -sv ${PM_DIR}/basecase_mul_pipe.v
|
||||
|
||||
# PolyMul zeta ROM
|
||||
xvlog -sv ${PM_DIR}/poly_mul_zeta_rom.v
|
||||
|
||||
144
sync_rtl/poly_mul/basecase_mul_pipe.v
Normal file
144
sync_rtl/poly_mul/basecase_mul_pipe.v
Normal file
@@ -0,0 +1,144 @@
|
||||
// basecase_mul_pipe.v - pipelined NTT-domain degree-1 multiplication
|
||||
//
|
||||
// Fixed-latency replacement for the combinational basecase_mul hot path.
|
||||
// The shared ntt/barrett_mul remains combinational for NTT users; this module
|
||||
// keeps the extra registers local to poly_mul_sync.
|
||||
|
||||
(* use_dsp = "no" *)
|
||||
module basecase_mul_pipe (
|
||||
input clk,
|
||||
input rst_n,
|
||||
input valid_i,
|
||||
input [11:0] a0, a1,
|
||||
input [11:0] b0, b1,
|
||||
input [11:0] zeta,
|
||||
output [11:0] c0,
|
||||
output [11:0] c1,
|
||||
output valid_o
|
||||
);
|
||||
localparam [11:0] Q12 = 12'd3329;
|
||||
localparam [12:0] K13 = 13'd5039;
|
||||
localparam [24:0] Q25 = 25'd3329;
|
||||
|
||||
function [12:0] barrett_q;
|
||||
input [23:0] p;
|
||||
reg [36:0] prod;
|
||||
begin
|
||||
prod = {13'd0, p} * K13;
|
||||
barrett_q = prod[36:24];
|
||||
end
|
||||
endfunction
|
||||
|
||||
function [11:0] barrett_reduce;
|
||||
input [23:0] p;
|
||||
input [12:0] qe;
|
||||
reg [24:0] q_approx;
|
||||
reg [24:0] r0;
|
||||
reg [24:0] r1;
|
||||
reg [24:0] r2;
|
||||
begin
|
||||
q_approx = qe * Q12;
|
||||
r0 = {1'b0, p} - q_approx;
|
||||
r1 = (r0 >= Q25) ? (r0 - Q25) : r0;
|
||||
r2 = (r1 >= Q25) ? (r1 - Q25) : r1;
|
||||
barrett_reduce = r2[11:0];
|
||||
end
|
||||
endfunction
|
||||
|
||||
function [11:0] mod_add;
|
||||
input [11:0] x;
|
||||
input [11:0] y;
|
||||
reg [12:0] sum;
|
||||
begin
|
||||
sum = {1'b0, x} + {1'b0, y};
|
||||
mod_add = (sum >= {1'b0, Q12}) ? (sum[11:0] - Q12) : sum[11:0];
|
||||
end
|
||||
endfunction
|
||||
|
||||
// Stage 1: first-level scalar products.
|
||||
reg [23:0] p10, p11, p12, p13;
|
||||
reg [11:0] zeta_s1;
|
||||
|
||||
// Stage 2: Barrett quotient estimates for first-level products.
|
||||
reg [23:0] p20, p21, p22, p23;
|
||||
reg [12:0] q20, q21, q22, q23;
|
||||
reg [11:0] zeta_s2;
|
||||
|
||||
// Stage 3: first-level reduced products.
|
||||
reg [11:0] t1_s3, t2_s3, t3_s3, t4_s3;
|
||||
reg [11:0] zeta_s3;
|
||||
|
||||
// Stage 4: zeta product for c0's second Barrett multiply.
|
||||
reg [11:0] t1_s4, t3_s4, t4_s4;
|
||||
reg [23:0] pz_s4;
|
||||
|
||||
// Stage 5: Barrett quotient estimate for zeta product.
|
||||
reg [11:0] t1_s5, t3_s5, t4_s5;
|
||||
reg [23:0] pz_s5;
|
||||
reg [12:0] qz_s5;
|
||||
|
||||
// Stage 6: reduced zeta product.
|
||||
reg [11:0] t1_s6, t3_s6, t4_s6, tz_s6;
|
||||
|
||||
// Stage 7: final modular additions.
|
||||
reg [11:0] c0_r, c1_r;
|
||||
reg [6:0] valid_sr;
|
||||
|
||||
assign c0 = c0_r;
|
||||
assign c1 = c1_r;
|
||||
assign valid_o = valid_sr[6];
|
||||
|
||||
always @(posedge clk or negedge rst_n) begin
|
||||
if (!rst_n) begin
|
||||
p10 <= 24'd0; p11 <= 24'd0; p12 <= 24'd0; p13 <= 24'd0; zeta_s1 <= 12'd0;
|
||||
p20 <= 24'd0; p21 <= 24'd0; p22 <= 24'd0; p23 <= 24'd0;
|
||||
q20 <= 13'd0; q21 <= 13'd0; q22 <= 13'd0; q23 <= 13'd0; zeta_s2 <= 12'd0;
|
||||
t1_s3 <= 12'd0; t2_s3 <= 12'd0; t3_s3 <= 12'd0; t4_s3 <= 12'd0; zeta_s3 <= 12'd0;
|
||||
t1_s4 <= 12'd0; t3_s4 <= 12'd0; t4_s4 <= 12'd0; pz_s4 <= 24'd0;
|
||||
t1_s5 <= 12'd0; t3_s5 <= 12'd0; t4_s5 <= 12'd0; pz_s5 <= 24'd0; qz_s5 <= 13'd0;
|
||||
t1_s6 <= 12'd0; t3_s6 <= 12'd0; t4_s6 <= 12'd0; tz_s6 <= 12'd0;
|
||||
c0_r <= 12'd0; c1_r <= 12'd0; valid_sr <= 7'd0;
|
||||
end else begin
|
||||
valid_sr <= {valid_sr[5:0], valid_i};
|
||||
|
||||
p10 <= {12'd0, a0} * b0;
|
||||
p11 <= {12'd0, a1} * b1;
|
||||
p12 <= {12'd0, a1} * b0;
|
||||
p13 <= {12'd0, a0} * b1;
|
||||
zeta_s1 <= zeta;
|
||||
|
||||
p20 <= p10; p21 <= p11; p22 <= p12; p23 <= p13;
|
||||
q20 <= barrett_q(p10);
|
||||
q21 <= barrett_q(p11);
|
||||
q22 <= barrett_q(p12);
|
||||
q23 <= barrett_q(p13);
|
||||
zeta_s2 <= zeta_s1;
|
||||
|
||||
t1_s3 <= barrett_reduce(p20, q20);
|
||||
t2_s3 <= barrett_reduce(p21, q21);
|
||||
t3_s3 <= barrett_reduce(p22, q22);
|
||||
t4_s3 <= barrett_reduce(p23, q23);
|
||||
zeta_s3 <= zeta_s2;
|
||||
|
||||
t1_s4 <= t1_s3;
|
||||
t3_s4 <= t3_s3;
|
||||
t4_s4 <= t4_s3;
|
||||
pz_s4 <= {12'd0, t2_s3} * zeta_s3;
|
||||
|
||||
t1_s5 <= t1_s4;
|
||||
t3_s5 <= t3_s4;
|
||||
t4_s5 <= t4_s4;
|
||||
pz_s5 <= pz_s4;
|
||||
qz_s5 <= barrett_q(pz_s4);
|
||||
|
||||
t1_s6 <= t1_s5;
|
||||
t3_s6 <= t3_s5;
|
||||
t4_s6 <= t4_s5;
|
||||
tz_s6 <= barrett_reduce(pz_s5, qz_s5);
|
||||
|
||||
c0_r <= mod_add(t1_s6, tz_s6);
|
||||
c1_r <= mod_add(t3_s6, t4_s6);
|
||||
end
|
||||
end
|
||||
|
||||
endmodule
|
||||
@@ -4,7 +4,7 @@
|
||||
// 256-coefficient NTT-domain polynomials.
|
||||
//
|
||||
// Operation flow:
|
||||
// IDLE → LOAD (256× A+B pairs) → COMP_CALC (read+compute, 1 cycle)
|
||||
// IDLE → LOAD (256× A+B pairs) → COMP_ISSUE → COMP_WAIT
|
||||
// → COMP_C0 (output c0) → COMP_C1 (output c1) → DONE → IDLE
|
||||
//
|
||||
// The LOAD phase accepts both A and B coefficients simultaneously
|
||||
@@ -39,10 +39,11 @@ module poly_mul_sync (
|
||||
// State definitions
|
||||
localparam S_IDLE = 3'd0;
|
||||
localparam S_LOAD = 3'd1;
|
||||
localparam S_COMP_CALC = 3'd2;
|
||||
localparam S_COMP_C0 = 3'd3;
|
||||
localparam S_COMP_C1 = 3'd4;
|
||||
localparam S_DONE = 3'd5;
|
||||
localparam S_COMP_ISSUE = 3'd2;
|
||||
localparam S_COMP_WAIT = 3'd3;
|
||||
localparam S_COMP_C0 = 3'd4;
|
||||
localparam S_COMP_C1 = 3'd5;
|
||||
localparam S_DONE = 3'd6;
|
||||
|
||||
reg [2:0] state, next_state;
|
||||
|
||||
@@ -72,16 +73,22 @@ module poly_mul_sync (
|
||||
.zeta (zeta)
|
||||
);
|
||||
|
||||
// Basecase multiply
|
||||
// Pipelined basecase multiply. One request is issued at a time; the
|
||||
// output interface is unchanged for top-level consumers.
|
||||
wire [11:0] bc_c0, bc_c1;
|
||||
basecase_mul u_bc (
|
||||
wire bc_vo;
|
||||
basecase_mul_pipe u_bc (
|
||||
.clk (clk),
|
||||
.rst_n(rst_n),
|
||||
.valid_i(state == S_COMP_ISSUE),
|
||||
.a0 (mem_a0),
|
||||
.a1 (mem_a1),
|
||||
.b0 (mem_b0),
|
||||
.b1 (mem_b1),
|
||||
.zeta(zeta),
|
||||
.c0 (bc_c0),
|
||||
.c1 (bc_c1)
|
||||
.c1 (bc_c1),
|
||||
.valid_o(bc_vo)
|
||||
);
|
||||
|
||||
// Output interface
|
||||
@@ -95,12 +102,13 @@ module poly_mul_sync (
|
||||
case (state)
|
||||
S_IDLE: if (valid_i && ready_o) next_state = S_LOAD;
|
||||
S_LOAD: if (load_cnt >= 255 && valid_i && ready_o)
|
||||
next_state = S_COMP_CALC;
|
||||
S_COMP_CALC: next_state = S_COMP_C0;
|
||||
next_state = S_COMP_ISSUE;
|
||||
S_COMP_ISSUE: next_state = S_COMP_WAIT;
|
||||
S_COMP_WAIT: if (bc_vo) next_state = S_COMP_C0;
|
||||
S_COMP_C0: if (valid_o && ready_i) next_state = S_COMP_C1;
|
||||
S_COMP_C1: if (valid_o && ready_i) begin
|
||||
if (comp_k >= 127) next_state = S_DONE;
|
||||
else next_state = S_COMP_CALC;
|
||||
else next_state = S_COMP_ISSUE;
|
||||
end
|
||||
S_DONE: next_state = S_IDLE;
|
||||
default: next_state = S_IDLE;
|
||||
@@ -139,8 +147,8 @@ module poly_mul_sync (
|
||||
end
|
||||
|
||||
// ---- COMPUTE phase ----
|
||||
// COMP_CALC: capture basecase_mul results
|
||||
if (state == S_COMP_CALC) begin
|
||||
// COMP_WAIT: capture pipelined basecase_mul results when ready.
|
||||
if (state == S_COMP_WAIT && bc_vo) begin
|
||||
c0_reg <= bc_c0;
|
||||
c1_reg <= bc_c1;
|
||||
end
|
||||
|
||||
@@ -17,6 +17,7 @@ xvlog -sv --relax -i . sync_rtl/ntt/zeta_rom.v
|
||||
xvlog -sv --relax -i . sync_rtl/ntt/butterfly_unit.v
|
||||
xvlog -sv --relax -i . sync_rtl/ntt/ntt_core.v
|
||||
xvlog -sv --relax -i . sync_rtl/poly_mul/basecase_mul.v
|
||||
xvlog -sv --relax -i . sync_rtl/poly_mul/basecase_mul_pipe.v
|
||||
xvlog -sv --relax -i . sync_rtl/poly_mul/poly_mul_zeta_rom.v
|
||||
xvlog -sv --relax -i . sync_rtl/poly_mul/poly_mul_sync.v
|
||||
xvlog -sv --relax -i . sync_rtl/storage/sd_bram.v
|
||||
|
||||
@@ -1009,7 +1009,7 @@ module mlkem_top #(
|
||||
// (d=du), then c2 = 1 poly of v (d=dv). Per poly = 256 coeffs -> 32*d
|
||||
// bytes (whole), so the bit buffer empties at each poly boundary.
|
||||
// micro-phase cp_ph: 0 present coeff addr; 1 feed comp_decomp (cd_valid);
|
||||
// 2 wait pipe; 3 capture compressed + accumulate bits; 4..n drain bytes.
|
||||
// 2 drop valid pulse; 3 wait/capture cd_vo; 4..n drain bytes.
|
||||
// ================================================================
|
||||
wire cd_active = (st == ST_ENC_C1) || (st == ST_ENC_C2);
|
||||
reg [11:0] cd_coeff; // coeff presented to comp_decomp
|
||||
@@ -1701,8 +1701,8 @@ module mlkem_top #(
|
||||
// ---- ST_ENC_C1/C2: Compress_d -> byteEncode_d -> ct region ----
|
||||
// C1 (E5): Compress_du(u[0..K-1]) from bank_se -> ct[0..c1_bytes).
|
||||
// C2 (E7): Compress_dv(v) from bank_t[UPSUM] -> ct[c1_bytes..ct_bytes).
|
||||
// Per coeff, 5-phase micro-sequence (read-ahead 1 cyc bram + 1 cyc
|
||||
// comp_decomp pipe), then a drain sub-phase emitting whole bytes:
|
||||
// Per coeff, 5-phase micro-sequence (read-ahead 1 cyc bram plus
|
||||
// variable-latency comp_decomp), then a drain sub-phase emitting whole bytes:
|
||||
// ph0: present coeff addr (cp_se_full / cp_bt_full by state).
|
||||
// ph1: coeff arrives (cp_coeff_src) -> latch into cd_coeff, pulse cd_valid.
|
||||
// ph2: drop cd_valid (1-cyc pulse); comp_decomp captures.
|
||||
@@ -1723,10 +1723,12 @@ module mlkem_top #(
|
||||
cp_ph <= 3'd3;
|
||||
end
|
||||
3'd3: begin
|
||||
// cd_out valid (cd_vo): append cp_d bits LSB-first at bit cp_nbits
|
||||
cp_buf <= cp_buf | (({13'd0, cd_out} & ((25'd1 << cp_d) - 25'd1)) << cp_nbits);
|
||||
cp_nbits <= cp_nbits + {1'b0, cp_d};
|
||||
cp_ph <= 3'd4;
|
||||
if (cd_vo) begin
|
||||
// Append compressed low cp_d bits LSB-first at bit cp_nbits.
|
||||
cp_buf <= cp_buf | (({13'd0, cd_out} & ((25'd1 << cp_d) - 25'd1)) << cp_nbits);
|
||||
cp_nbits <= cp_nbits + {1'b0, cp_d};
|
||||
cp_ph <= 3'd4;
|
||||
end
|
||||
end
|
||||
default: begin // 3'd4: drain whole bytes
|
||||
if (cp_nbits >= 6'd8) begin
|
||||
|
||||
Reference in New Issue
Block a user