fallensigh/ciciec2026
Archived
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf(mlkem): dual-butterfly NTT core (R1) — ~12% fewer Encaps cycles

Browse Source 
Issue 2 butterflies/cycle from the same NTT group (positions p,p+1).
ML-KEM's 7 layers have layer_len in {128..2} (always even, never 1),
so each pair shares one zeta and touches 4 distinct coefficient indices
— no intra-group RAW conflict, single zeta ROM. Compute per transform
~966->~490 cycles.

Verified: standalone ntt TB 19/19 (fwd+inv), full KAT Encaps 9/9 +
Decaps 9/9 + KeyGen + hello_world all PASS (all K).
This commit is contained in:
FallenSigh
2026-07-11 03:05:35 +08:00
parent 975903000e
commit 454c3e559d
1 changed files with 75 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
rtl/ip/mlkem/ntt/ntt_core.v
View File

@@ -1,9 +1,13 @@
// ntt_core.v - NTT core with individual coefficient registers // ntt_core.v - NTT core with individual coefficient registers (2 butterflies)
// //
// Uses 256 individual 12-bit registers and a deeply pipelined butterfly path. // Uses 256 individual 12-bit registers and TWO deeply pipelined butterfly paths
// The arithmetic hot path is split into: // running in lock-step (R1: dual-butterfly parallelism). Each issue cycle emits
// address/operand/zeta register -> pipelined Barrett butterfly -> writeback // two butterflies from the SAME group at positions (p, p+1); since ML-KEM's NTT
// Each NTT layer is issued continuously into the butterfly pipeline, then the // has 7 layers with layer_len in {128..2} (always even, never 1), the pair always
// shares one zeta and touches four distinct coefficient indices, so a single
// zeta ROM suffices and there is no intra-group read/write conflict.
// address/operand/zeta register -> 2x pipelined Barrett butterfly -> writeback
// Each NTT layer issues 2 butterflies/cycle (64 cycles/layer vs 128), then the
// core drains pending writebacks before starting the next dependent layer. // core drains pending writebacks before starting the next dependent layer.
// In inverse mode, final x3303 output scaling also uses a pipelined Barrett // In inverse mode, final x3303 output scaling also uses a pipelined Barrett
// multiplier so the output path does not reintroduce a combinational reducer. // multiplier so the output path does not reintroduce a combinational reducer.
@@ -45,13 +49,16 @@ module ntt_core (
reg bf_done; reg bf_done;
reg mode_r; reg mode_r;
reg layer_issue_done; reg layer_issue_done;
reg [4:0] bf_inflight; reg [5:0] bf_inflight; // pending issue cycles (pairs) in flight
reg [7:0] wa_d0, wa_d1, wa_d2, wa_d3, wa_d4, wa_d5, wa_d6, wa_d7; // Two writeback address delay lines (lane0: j, lane1: j+1), each +len.
reg [7:0] wb_d0, wb_d1, wb_d2, wb_d3, wb_d4, wb_d5, wb_d6, wb_d7; reg [7:0] a0_d0,a0_d1,a0_d2,a0_d3,a0_d4,a0_d5,a0_d6,a0_d7;
reg [7:0] b0_d0,b0_d1,b0_d2,b0_d3,b0_d4,b0_d5,b0_d6,b0_d7;
reg [7:0] a1_d0,a1_d1,a1_d2,a1_d3,a1_d4,a1_d5,a1_d6,a1_d7;
reg [7:0] b1_d0,b1_d1,b1_d2,b1_d3,b1_d4,b1_d5,b1_d6,b1_d7;
reg wb_valid_r; reg wb_valid_r;
reg [DW-1:0] wb_a_data_r, wb_b_data_r; reg [DW-1:0] wb_a0_r, wb_b0_r, wb_a1_r, wb_b1_r;
reg [7:0] wb_wa_r, wb_wb_r; reg [7:0] wb_a0a_r, wb_b0a_r, wb_a1a_r, wb_b1a_r;
reg [DW-1:0] coeff_out_r; reg [DW-1:0] coeff_out_r;
reg valid_o_r; reg valid_o_r;
@@ -63,38 +70,42 @@ module ntt_core (
wire [8:0] group_end = {1'b0, start} + {1'b0, layer_len}; wire [8:0] group_end = {1'b0, start} + {1'b0, layer_len};
wire [8:0] next_group_start = {1'b0, start} + {1'b0, layer_len} + {1'b0, layer_len}; wire [8:0] next_group_start = {1'b0, start} + {1'b0, layer_len} + {1'b0, layer_len};
wire issue_group_last = ({1'b0, j} + 9'd1 >= group_end); // pair (j, j+1) is the last in this group when j+2 reaches group_end
wire issue_group_last = ({1'b0, j} + 9'd2 >= group_end);
wire issue_layer_last = issue_group_last && (next_group_start >= 9'd256); wire issue_layer_last = issue_group_last && (next_group_start >= 9'd256);
wire [7:0] next_layer_len = mode_r ? (layer_len << 1) : (layer_len >> 1); wire [7:0] next_layer_len = mode_r ? (layer_len << 1) : (layer_len >> 1);
wire issue_fire = (state == S_CMP_RUN) && !layer_issue_done; wire issue_fire = (state == S_CMP_RUN) && !layer_issue_done;
wire [DW-1:0] issue_a = cr[j];
wire [DW-1:0] issue_b = cr[j + layer_len];
wire [DW-1:0] bf_a_out, bf_b_out; // lane0 = position j, lane1 = position j+1 (same group => same zeta)
wire bf_valid; wire [DW-1:0] issue_a0 = cr[j];
butterfly_unit_pipe u_bf ( wire [DW-1:0] issue_b0 = cr[j + layer_len];
.clk(clk), wire [DW-1:0] issue_a1 = cr[j + 8'd1];
.rst_n(rst_n), wire [DW-1:0] issue_b1 = cr[j + 8'd1 + layer_len];
wire [DW-1:0] bf0_a_out, bf0_b_out, bf1_a_out, bf1_b_out;
wire bf0_valid, bf1_valid;
butterfly_unit_pipe u_bf0 (
.clk(clk), .rst_n(rst_n),
.valid_i(issue_fire), .valid_i(issue_fire),
.a(issue_a), .a(issue_a0), .b(issue_b0),
.b(issue_b), .zeta(zeta), .mode(mode_r),
.zeta(zeta), .a_out(bf0_a_out), .b_out(bf0_b_out), .valid_o(bf0_valid)
.mode(mode_r), );
.a_out(bf_a_out), butterfly_unit_pipe u_bf1 (
.b_out(bf_b_out), .clk(clk), .rst_n(rst_n),
.valid_o(bf_valid) .valid_i(issue_fire),
.a(issue_a1), .b(issue_b1),
.zeta(zeta), .mode(mode_r),
.a_out(bf1_a_out), .b_out(bf1_b_out), .valid_o(bf1_valid)
); );
wire [DW-1:0] scale_product; wire [DW-1:0] scale_product;
wire scale_valid_o; wire scale_valid_o;
barrett_mul_pipe u_scl ( barrett_mul_pipe u_scl (
.clk(clk), .clk(clk), .rst_n(rst_n),
.rst_n(rst_n),
.valid_i(scale_valid_i), .valid_i(scale_valid_i),
.a(scale_a_i), .a(scale_a_i), .b(12'd3303),
.b(12'd3303), .product(scale_product), .valid_o(scale_valid_o)
.product(scale_product),
.valid_o(scale_valid_o)
); );
assign ready_o = (state == S_IDLE) || (state == S_LOAD); assign ready_o = (state == S_IDLE) || (state == S_LOAD);
@@ -108,7 +119,7 @@ module ntt_core (
S_IDLE: if (valid_i) next_state = S_LOAD; S_IDLE: if (valid_i) next_state = S_LOAD;
S_LOAD: if (load_cnt >= 8'd255 && valid_i) next_state = S_CMP_RUN; S_LOAD: if (load_cnt >= 8'd255 && valid_i) next_state = S_CMP_RUN;
S_CMP_RUN: if (layer_issue_done) next_state = S_CMP_DRAIN; S_CMP_RUN: if (layer_issue_done) next_state = S_CMP_DRAIN;
S_CMP_DRAIN: if (bf_inflight == 5'd0) S_CMP_DRAIN: if (bf_inflight == 6'd0)
next_state = (layer + 3'd1 >= LAYERS) ? S_OUT_PREP : S_CMP_RUN; next_state = (layer + 3'd1 >= LAYERS) ? S_OUT_PREP : S_CMP_RUN;
S_OUT_PREP: next_state = mode_r ? S_OUT_SCALE : S_OUTPUT; 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_OUTPUT: if (valid_o_r && ready_i && out_cnt >= 8'd255) next_state = S_DONE;
@@ -117,7 +128,6 @@ module ntt_core (
default: next_state = S_IDLE; default: next_state = S_IDLE;
endcase endcase
end end
always @(posedge clk or negedge rst_n) begin always @(posedge clk or negedge rst_n) begin
if (!rst_n) begin if (!rst_n) begin
state <= S_IDLE; state <= S_IDLE;
@@ -133,16 +143,14 @@ module ntt_core (
bf_done <= 1'b0; bf_done <= 1'b0;
mode_r <= 1'b0; mode_r <= 1'b0;
layer_issue_done <= 1'b0; layer_issue_done <= 1'b0;
bf_inflight <= 5'd0; bf_inflight <= 6'd0;
wa_d0 <= 8'd0; wa_d1 <= 8'd0; wa_d2 <= 8'd0; wa_d3 <= 8'd0; a0_d0<=8'd0;a0_d1<=8'd0;a0_d2<=8'd0;a0_d3<=8'd0;a0_d4<=8'd0;a0_d5<=8'd0;a0_d6<=8'd0;a0_d7<=8'd0;
wa_d4 <= 8'd0; wa_d5 <= 8'd0; wa_d6 <= 8'd0; wa_d7 <= 8'd0; b0_d0<=8'd0;b0_d1<=8'd0;b0_d2<=8'd0;b0_d3<=8'd0;b0_d4<=8'd0;b0_d5<=8'd0;b0_d6<=8'd0;b0_d7<=8'd0;
wb_d0 <= 8'd0; wb_d1 <= 8'd0; wb_d2 <= 8'd0; wb_d3 <= 8'd0; a1_d0<=8'd0;a1_d1<=8'd0;a1_d2<=8'd0;a1_d3<=8'd0;a1_d4<=8'd0;a1_d5<=8'd0;a1_d6<=8'd0;a1_d7<=8'd0;
wb_d4 <= 8'd0; wb_d5 <= 8'd0; wb_d6 <= 8'd0; wb_d7 <= 8'd0; b1_d0<=8'd0;b1_d1<=8'd0;b1_d2<=8'd0;b1_d3<=8'd0;b1_d4<=8'd0;b1_d5<=8'd0;b1_d6<=8'd0;b1_d7<=8'd0;
wb_valid_r <= 1'b0; wb_valid_r <= 1'b0;
wb_a_data_r <= 12'd0; wb_a0_r <= 12'd0; wb_b0_r <= 12'd0; wb_a1_r <= 12'd0; wb_b1_r <= 12'd0;
wb_b_data_r <= 12'd0; wb_a0a_r <= 8'd0; wb_b0a_r <= 8'd0; wb_a1a_r <= 8'd0; wb_b1a_r <= 8'd0;
wb_wa_r <= 8'd0;
wb_wb_r <= 8'd0;
coeff_out_r <= 12'd0; coeff_out_r <= 12'd0;
valid_o_r <= 1'b0; valid_o_r <= 1'b0;
scale_valid_i <= 1'b0; scale_valid_i <= 1'b0;
@@ -152,15 +160,17 @@ module ntt_core (
state <= next_state; state <= next_state;
scale_valid_i <= 1'b0; scale_valid_i <= 1'b0;
wa_d0 <= j; wa_d1 <= wa_d0; wa_d2 <= wa_d1; wa_d3 <= wa_d2; // Writeback address delay lines (8-deep = butterfly latency).
wa_d4 <= wa_d3; wa_d5 <= wa_d4; wa_d6 <= wa_d5; wa_d7 <= wa_d6; a0_d0<=j; a0_d1<=a0_d0;a0_d2<=a0_d1;a0_d3<=a0_d2;a0_d4<=a0_d3;a0_d5<=a0_d4;a0_d6<=a0_d5;a0_d7<=a0_d6;
wb_d0 <= j + layer_len; wb_d1 <= wb_d0; wb_d2 <= wb_d1; wb_d3 <= wb_d2; b0_d0<=j+layer_len; b0_d1<=b0_d0;b0_d2<=b0_d1;b0_d3<=b0_d2;b0_d4<=b0_d3;b0_d5<=b0_d4;b0_d6<=b0_d5;b0_d7<=b0_d6;
wb_d4 <= wb_d3; wb_d5 <= wb_d4; wb_d6 <= wb_d5; wb_d7 <= wb_d6; a1_d0<=j+8'd1; a1_d1<=a1_d0;a1_d2<=a1_d1;a1_d3<=a1_d2;a1_d4<=a1_d3;a1_d5<=a1_d4;a1_d6<=a1_d5;a1_d7<=a1_d6;
b1_d0<=j+8'd1+layer_len; b1_d1<=b1_d0;b1_d2<=b1_d1;b1_d3<=b1_d2;b1_d4<=b1_d3;b1_d5<=b1_d4;b1_d6<=b1_d5;b1_d7<=b1_d6;
// In-flight pair accounting (mirrors single-butterfly version).
if (issue_fire && !wb_valid_r) if (issue_fire && !wb_valid_r)
bf_inflight <= bf_inflight + 5'd1; bf_inflight <= bf_inflight + 6'd1;
else if (!issue_fire && wb_valid_r) else if (!issue_fire && wb_valid_r)
bf_inflight <= bf_inflight - 5'd1; bf_inflight <= bf_inflight - 6'd1;
if (state != S_OUTPUT && state != S_OUT_SCALE) if (state != S_OUTPUT && state != S_OUT_SCALE)
valid_o_r <= 1'b0; valid_o_r <= 1'b0;
@@ -176,7 +186,7 @@ module ntt_core (
layer <= 3'd0; layer <= 3'd0;
bf_done <= 1'b0; bf_done <= 1'b0;
layer_issue_done <= 1'b0; layer_issue_done <= 1'b0;
bf_inflight <= 5'd0; bf_inflight <= 6'd0;
mode_r <= mode; mode_r <= mode;
if (!mode) begin if (!mode) begin
layer_len <= 8'd128; layer_len <= 8'd128;
@@ -192,6 +202,7 @@ module ntt_core (
load_cnt <= load_cnt + 8'd1; load_cnt <= load_cnt + 8'd1;
end end
// ---- issue advance: step j by 2 (two butterflies/cycle) ----
if (state == S_CMP_RUN && !layer_issue_done) begin if (state == S_CMP_RUN && !layer_issue_done) begin
if (issue_group_last) begin if (issue_group_last) begin
if (!mode_r) zeta_idx <= zeta_idx + 7'd1; if (!mode_r) zeta_idx <= zeta_idx + 7'd1;
@@ -204,24 +215,28 @@ module ntt_core (
j <= next_group_start[7:0]; j <= next_group_start[7:0];
end end
end else begin end else begin
j <= j + 8'd1; j <= j + 8'd2;
end end
end end
// ---- writeback: commit previous cycle's pair results ----
if (wb_valid_r) begin if (wb_valid_r) begin
cr[wb_wa_r] <= wb_a_data_r; cr[wb_a0a_r] <= wb_a0_r;
cr[wb_wb_r] <= wb_b_data_r; cr[wb_b0a_r] <= wb_b0_r;
cr[wb_a1a_r] <= wb_a1_r;
cr[wb_b1a_r] <= wb_b1_r;
end end
wb_valid_r <= bf_valid; wb_valid_r <= bf0_valid; // both lanes share latency/valid
if (bf_valid) begin if (bf0_valid) begin
wb_a_data_r <= bf_a_out; wb_a0_r <= bf0_a_out; wb_b0_r <= bf0_b_out;
wb_b_data_r <= bf_b_out; wb_a1_r <= bf1_a_out; wb_b1_r <= bf1_b_out;
wb_wa_r <= wa_d7; wb_a0a_r <= a0_d7; wb_b0a_r <= b0_d7;
wb_wb_r <= wb_d7; wb_a1a_r <= a1_d7; wb_b1a_r <= b1_d7;
end end
if (state == S_CMP_DRAIN && bf_inflight == 5'd0) begin // ---- layer advance after drain ----
if (state == S_CMP_DRAIN && bf_inflight == 6'd0) begin
if (layer + 3'd1 >= LAYERS) begin if (layer + 3'd1 >= LAYERS) begin
bf_done <= 1'b1; bf_done <= 1'b1;
end else begin end else begin
@@ -260,7 +275,6 @@ module ntt_core (
scale_a_i <= cr[scale_issue_cnt[7:0]]; scale_a_i <= cr[scale_issue_cnt[7:0]];
scale_issue_cnt <= scale_issue_cnt + 9'd1; scale_issue_cnt <= scale_issue_cnt + 9'd1;
end end
valid_o_r <= 1'b0; valid_o_r <= 1'b0;
if (scale_valid_o) begin if (scale_valid_o) begin
coeff_out_r <= scale_product; coeff_out_r <= scale_product;
@@ -275,11 +289,10 @@ module ntt_core (
scale_issue_cnt <= 9'd0; scale_issue_cnt <= 9'd0;
scale_emit_cnt <= 9'd0; scale_emit_cnt <= 9'd0;
layer_issue_done <= 1'b0; layer_issue_done <= 1'b0;
bf_inflight <= 5'd0; bf_inflight <= 6'd0;
wb_valid_r <= 1'b0; wb_valid_r <= 1'b0;
valid_o_r <= 1'b0; valid_o_r <= 1'b0;
end end
end end
end end
endmodule endmodule