fallensigh/mlkem-sync
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Pipeline ML-KEM datapath bottlenecks

Browse Source
This commit is contained in:
FallenSigh
2026-07-08 00:23:46 +08:00
parent 8c3f4317f5
commit 372a90e601
17 changed files with 776 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

232
sync_rtl/ntt/ntt_core.v
View File

@@ -1,8 +1,10 @@
// ntt_core.v - NTT core with individual coefficient registers
//
// Uses 256 individual 12-bit registers and generate-based muxing
// to avoid any part-select simulation issues.
// 3-cycle butterfly: SetAddr -> Read -> Compute+Write
// Uses 256 individual 12-bit registers and a deeply pipelined butterfly path.
// The arithmetic hot path is split into:
// address -> operand/zeta register -> pipelined Barrett butterfly -> writeback
// In inverse mode, final x3303 output scaling also uses a pipelined Barrett
// multiplier so the output path does not reintroduce a combinational reducer.
module ntt_core (
input clk, rst_n,
@@ -17,72 +19,149 @@ module ntt_core (
);
localparam N = 256, LAYERS = 7, DW = 12;
// Individual coefficient registers
reg [DW-1:0] cr [0:N-1];
integer ci;
// State machine
localparam S_IDLE=3'd0, S_LOAD=3'd1, S_CMP_A=3'd2, S_CMP_B=3'd3,
S_CMP_C=3'd4, S_OUTPUT=3'd5, S_DONE=3'd6;
reg [2:0] state, next_state;
localparam S_IDLE = 4'd0;
localparam S_LOAD = 4'd1;
localparam S_CMP_A = 4'd2;
localparam S_CMP_B = 4'd3;
localparam S_CMP_ISSUE = 4'd4;
localparam S_CMP_WAIT = 4'd5;
localparam S_CMP_WB = 4'd6;
localparam S_OUT_PREP = 4'd7;
localparam S_OUTPUT = 4'd8;
localparam S_OUT_SCALE = 4'd9;
localparam S_DONE = 4'd10;
reg [7:0] load_cnt, out_cnt;
reg [3:0] state, next_state;
reg [7:0] load_cnt;
reg [7:0] out_cnt;
reg [8:0] scale_issue_cnt;
reg [8:0] scale_emit_cnt;
reg [7:0] j, start, layer_len;
reg [6:0] zeta_idx;
reg [2:0] layer;
reg bf_done;
reg bf_done;
reg mode_r;
// Pipeline registers
reg [DW-1:0] r_a, r_b;
reg [7:0] r_wa, r_wb;
reg [DW-1:0] r_a, r_b, r_zeta;
reg [7:0] r_wa, r_wb;
reg [DW-1:0] wr_a_data, wr_b_data;
reg [7:0] wr_wa, wr_wb;
reg [DW-1:0] coeff_out_r;
reg valid_o_r;
reg scale_valid_i;
reg [DW-1:0] scale_a_i;
// Zeta
wire [DW-1:0] zeta;
zeta_rom u_z (.addr(zeta_idx), .zeta(zeta));
// Butterfly
wire [DW-1:0] bf_a_out, bf_b_out;
butterfly_unit u_bf (
.a(r_a), .b(r_b), .zeta(zeta), .mode(mode),
.a_out(bf_a_out), .b_out(bf_b_out));
wire bf_valid;
butterfly_unit_pipe u_bf (
.clk(clk),
.rst_n(rst_n),
.valid_i(state == S_CMP_ISSUE),
.a(r_a),
.b(r_b),
.zeta(r_zeta),
.mode(mode_r),
.a_out(bf_a_out),
.b_out(bf_b_out),
.valid_o(bf_valid)
);
// Output scaling
wire [DW-1:0] coeff_scaled;
barrett_mul u_scl (.a(cr[out_cnt]), .b(12'd3303), .product(coeff_scaled));
assign coeff_out = mode ? coeff_scaled : cr[out_cnt];
wire [DW-1:0] scale_product;
wire scale_valid_o;
barrett_mul_pipe u_scl (
.clk(clk),
.rst_n(rst_n),
.valid_i(scale_valid_i),
.a(scale_a_i),
.b(12'd3303),
.product(scale_product),
.valid_o(scale_valid_o)
);
assign ready_o = (state == S_IDLE) || (state == S_LOAD);
assign valid_o = (state == S_OUTPUT);
assign done_o = (state == S_DONE);
assign ready_o = (state == S_IDLE) || (state == S_LOAD);
assign coeff_out = coeff_out_r;
assign valid_o = valid_o_r;
assign done_o = (state == S_DONE);
always @* begin
next_state = state;
case (state)
S_IDLE: if (valid_i) next_state = S_LOAD;
S_LOAD: if (load_cnt >= 255 && valid_i) next_state = S_CMP_A;
S_CMP_A: if (bf_done) next_state = S_OUTPUT; else next_state = S_CMP_B;
S_CMP_B: if (bf_done) next_state = S_OUTPUT; else next_state = S_CMP_C;
S_CMP_C: if (bf_done) next_state = S_OUTPUT; else next_state = S_CMP_A;
S_OUTPUT:if (out_cnt >= 255 && ready_i) next_state = S_DONE;
S_DONE: next_state = S_IDLE;
default: next_state = S_IDLE;
S_IDLE: if (valid_i) next_state = S_LOAD;
S_LOAD: if (load_cnt >= 8'd255 && valid_i) next_state = S_CMP_A;
S_CMP_A: next_state = bf_done ? S_OUT_PREP : S_CMP_B;
S_CMP_B: next_state = bf_done ? S_OUT_PREP : S_CMP_ISSUE;
S_CMP_ISSUE: next_state = S_CMP_WAIT;
S_CMP_WAIT: if (bf_valid) next_state = S_CMP_WB;
S_CMP_WB: next_state = S_CMP_A;
S_OUT_PREP: next_state = mode_r ? S_OUT_SCALE : S_OUTPUT;
S_OUTPUT: if (valid_o_r && ready_i && out_cnt >= 8'd255) next_state = S_DONE;
S_OUT_SCALE: if (scale_valid_o && scale_emit_cnt >= 9'd255) next_state = S_DONE;
S_DONE: next_state = S_IDLE;
default: next_state = S_IDLE;
endcase
end
always @(posedge clk or negedge rst_n) begin
if (!rst_n) begin
state<=S_IDLE; load_cnt<=0; out_cnt<=0; j<=0; start<=0; layer_len<=0;
zeta_idx<=0; layer<=0; bf_done<=0; r_a<=0; r_b<=0; r_wa<=0; r_wb<=0;
for (ci=0; ci<N; ci=ci+1) cr[ci] <= 0;
state <= S_IDLE;
load_cnt <= 8'd0;
out_cnt <= 8'd0;
scale_issue_cnt <= 9'd0;
scale_emit_cnt <= 9'd0;
j <= 8'd0;
start <= 8'd0;
layer_len <= 8'd0;
zeta_idx <= 7'd0;
layer <= 3'd0;
bf_done <= 1'b0;
mode_r <= 1'b0;
r_a <= 12'd0;
r_b <= 12'd0;
r_zeta <= 12'd0;
r_wa <= 8'd0;
r_wb <= 8'd0;
wr_a_data <= 12'd0;
wr_b_data <= 12'd0;
wr_wa <= 8'd0;
wr_wb <= 8'd0;
coeff_out_r <= 12'd0;
valid_o_r <= 1'b0;
scale_valid_i <= 1'b0;
scale_a_i <= 12'd0;
for (ci = 0; ci < N; ci = ci + 1) cr[ci] <= 12'd0;
end else begin
state <= next_state;
scale_valid_i <= 1'b0;
if (state != S_OUTPUT && state != S_OUT_SCALE)
valid_o_r <= 1'b0;
if (state == S_IDLE && valid_i) begin
cr[0] <= coeff_in;
load_cnt<=1; out_cnt<=0; j<=0; start<=0; layer<=0; bf_done<=0;
if (!mode) begin layer_len<=128; zeta_idx<=1; end
else begin layer_len<=2; zeta_idx<=127; end
load_cnt <= 8'd1;
out_cnt <= 8'd0;
scale_issue_cnt <= 9'd0;
scale_emit_cnt <= 9'd0;
j <= 8'd0;
start <= 8'd0;
layer <= 3'd0;
bf_done <= 1'b0;
mode_r <= mode;
if (!mode) begin
layer_len <= 8'd128;
zeta_idx <= 7'd1;
end else begin
layer_len <= 8'd2;
zeta_idx <= 7'd127;
end
end
if (state == S_LOAD && valid_i) begin
@@ -90,32 +169,38 @@ module ntt_core (
load_cnt <= load_cnt + 8'd1;
end
// S_CMP_A: set read addresses (j, j+len)
if (state == S_CMP_A) begin
r_wa <= j;
r_wb <= j + layer_len;
end
// S_CMP_B: capture read data
if (state == S_CMP_B) begin
r_a <= cr[j];
r_b <= cr[j + layer_len];
r_zeta <= zeta;
end
// S_CMP_C: write butterfly results, advance counters
if (state == S_CMP_C) begin
cr[r_wa] <= bf_a_out;
cr[r_wb] <= bf_b_out;
if (state == S_CMP_WAIT && bf_valid) begin
wr_a_data <= bf_a_out;
wr_b_data <= bf_b_out;
wr_wa <= r_wa;
wr_wb <= r_wb;
end
if (state == S_CMP_WB) begin
cr[wr_wa] <= wr_a_data;
cr[wr_wb] <= wr_b_data;
j <= j + 8'd1;
if (j + 8'd1 >= start + layer_len) begin
if (!mode) zeta_idx <= zeta_idx + 7'd1;
else zeta_idx <= zeta_idx - 7'd1;
if (!mode_r) zeta_idx <= zeta_idx + 7'd1;
else zeta_idx <= zeta_idx - 7'd1;
if ({1'b0,start} + {1'b0,layer_len} + {1'b0,layer_len} >= 256) begin
if ({1'b0,start} + {1'b0,layer_len} + {1'b0,layer_len} >= 9'd256) begin
layer <= layer + 3'd1;
layer_len <= mode ? (layer_len<<1) : (layer_len>>1);
start <= 0; j <= 0;
layer_len <= mode_r ? (layer_len << 1) : (layer_len >> 1);
start <= 8'd0;
j <= 8'd0;
if (layer + 3'd1 >= LAYERS) bf_done <= 1'b1;
end else begin
start <= start + layer_len + layer_len;
@@ -124,8 +209,49 @@ module ntt_core (
end
end
if (state == S_OUTPUT && ready_i)
out_cnt <= (out_cnt>=255) ? 0 : (out_cnt+8'd1);
if (state == S_OUT_PREP) begin
out_cnt <= 8'd0;
scale_issue_cnt <= 9'd0;
scale_emit_cnt <= 9'd0;
if (!mode_r) begin
coeff_out_r <= cr[0];
valid_o_r <= 1'b1;
end
end
if (state == S_OUTPUT && valid_o_r && ready_i) begin
if (out_cnt < 8'd255) begin
out_cnt <= out_cnt + 8'd1;
coeff_out_r <= cr[out_cnt + 8'd1];
valid_o_r <= 1'b1;
end else begin
out_cnt <= 8'd0;
valid_o_r <= 1'b0;
end
end
if (state == S_OUT_SCALE) begin
if (scale_issue_cnt < 9'd256) begin
scale_valid_i <= 1'b1;
scale_a_i <= cr[scale_issue_cnt[7:0]];
scale_issue_cnt <= scale_issue_cnt + 9'd1;
end
valid_o_r <= 1'b0;
if (scale_valid_o) begin
coeff_out_r <= scale_product;
valid_o_r <= 1'b1;
scale_emit_cnt <= scale_emit_cnt + 9'd1;
end
end
if (state == S_DONE) begin
load_cnt <= 8'd0;
out_cnt <= 8'd0;
scale_issue_cnt <= 9'd0;
scale_emit_cnt <= 9'd0;
valid_o_r <= 1'b0;
end
end
end