From 4538ec496ff1d2385fb85e35f8e7e6ae96430320 Mon Sep 17 00:00:00 2001 From: FallenSigh Date: Sat, 11 Jul 2026 07:05:12 +0800 Subject: [PATCH] perf(mlkem): register butterfly inputs to cut the 100 MHz NTT critical path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit At mlkem_clk=100 MHz the worst path (WNS -2.14 ns) was coeff_bram read -> butterfly diff (b-a) -> barrett first multiply stage, all in one cycle. The butterfly took its operands combinationally from the BRAM read (unlike poly_mul's basecase, which already registers them). Add an input register stage (a/b/zeta/mode/valid) in butterfly_unit_pipe so the BRAM read and the first multiply land in separate cycles; feed the a/b/mode alignment delay lines from the registered inputs. Butterfly latency 8->9; ntt_core writeback address delay lines deepened 9->10 to match. Verified byte-exact vs NIST KAT: NTT standalone TB 19/19 (fwd+inv), Encaps 9/9, Decaps 9/9, hello_world e2e — all PASS. Co-Authored-By: Claude Opus 4.8 (1M context) --- rtl/ip/mlkem/ntt/butterfly_unit_pipe.v | 29 +++++++++++++++++++------- rtl/ip/mlkem/ntt/ntt_core.v | 16 +++++++------- 2 files changed, 30 insertions(+), 15 deletions(-) diff --git a/rtl/ip/mlkem/ntt/butterfly_unit_pipe.v b/rtl/ip/mlkem/ntt/butterfly_unit_pipe.v index c672953..c365a24 100644 --- a/rtl/ip/mlkem/ntt/butterfly_unit_pipe.v +++ b/rtl/ip/mlkem/ntt/butterfly_unit_pipe.v @@ -18,21 +18,32 @@ module butterfly_unit_pipe ( localparam [11:0] Q12 = 12'd3329; localparam [12:0] Q13 = 13'd3329; - wire [11:0] diff = (b >= a) ? (b - a) : (b - a + Q12); - wire [11:0] mul_data = mode ? diff : b; + // Plan A: register the inputs. Operands arrive combinationally from the + // coeff_bram read; without this stage the path coeff_bram-read -> diff -> + // first multiply stage sits in one cycle (the mlkem_clk critical path at + // 100 MHz). Registering here splits it into (BRAM read -> reg) and + // (reg -> diff -> multiply). Adds one pipeline cycle (absorbed by the NTT + // writeback delay line). + reg [11:0] a_r, b_r, zeta_r; + reg mode_ri, vld_r; + + wire [11:0] diff = (b_r >= a_r) ? (b_r - a_r) : (b_r - a_r + Q12); + wire [11:0] mul_data = mode_ri ? diff : b_r; wire [11:0] mul_result; wire mul_valid; barrett_mul_pipe u_barrett ( .clk(clk), .rst_n(rst_n), - .valid_i(valid_i), - .a(zeta), + .valid_i(vld_r), + .a(zeta_r), .b(mul_data), .product(mul_result), .valid_o(mul_valid) ); + // a/b/mode delay lines (fed from the REGISTERED inputs a_r/b_r/mode_ri so + // a_d6 aligns with mul_valid = valid_i delayed 8). reg [11:0] a_d0, a_d1, a_d2, a_d3, a_d4, a_d5, a_d6; reg [11:0] b_d0, b_d1, b_d2, b_d3, b_d4, b_d5, b_d6; reg mode_d0, mode_d1, mode_d2, mode_d3, mode_d4, mode_d5, mode_d6; @@ -51,6 +62,7 @@ module butterfly_unit_pipe ( always @(posedge clk or negedge rst_n) begin if (!rst_n) begin + a_r <= 12'd0; b_r <= 12'd0; zeta_r <= 12'd0; mode_ri <= 1'b0; vld_r <= 1'b0; a_d0 <= 12'd0; a_d1 <= 12'd0; a_d2 <= 12'd0; a_d3 <= 12'd0; a_d4 <= 12'd0; a_d5 <= 12'd0; a_d6 <= 12'd0; b_d0 <= 12'd0; b_d1 <= 12'd0; b_d2 <= 12'd0; b_d3 <= 12'd0; b_d4 <= 12'd0; b_d5 <= 12'd0; b_d6 <= 12'd0; mode_d0 <= 1'b0; mode_d1 <= 1'b0; mode_d2 <= 1'b0; mode_d3 <= 1'b0; mode_d4 <= 1'b0; mode_d5 <= 1'b0; mode_d6 <= 1'b0; @@ -58,9 +70,12 @@ module butterfly_unit_pipe ( b_out_r <= 12'd0; valid_r <= 1'b0; end else begin - a_d0 <= a; a_d1 <= a_d0; a_d2 <= a_d1; a_d3 <= a_d2; a_d4 <= a_d3; a_d5 <= a_d4; a_d6 <= a_d5; - b_d0 <= b; b_d1 <= b_d0; b_d2 <= b_d1; b_d3 <= b_d2; b_d4 <= b_d3; b_d5 <= b_d4; b_d6 <= b_d5; - mode_d0 <= mode; mode_d1 <= mode_d0; mode_d2 <= mode_d1; + // input register stage + a_r <= a; b_r <= b; zeta_r <= zeta; mode_ri <= mode; vld_r <= valid_i; + + a_d0 <= a_r; a_d1 <= a_d0; a_d2 <= a_d1; a_d3 <= a_d2; a_d4 <= a_d3; a_d5 <= a_d4; a_d6 <= a_d5; + b_d0 <= b_r; b_d1 <= b_d0; b_d2 <= b_d1; b_d3 <= b_d2; b_d4 <= b_d3; b_d5 <= b_d4; b_d6 <= b_d5; + mode_d0 <= mode_ri; mode_d1 <= mode_d0; mode_d2 <= mode_d1; mode_d3 <= mode_d2; mode_d4 <= mode_d3; mode_d5 <= mode_d4; mode_d6 <= mode_d5; valid_r <= mul_valid; diff --git a/rtl/ip/mlkem/ntt/ntt_core.v b/rtl/ip/mlkem/ntt/ntt_core.v index 681c1b8..8e46be2 100644 --- a/rtl/ip/mlkem/ntt/ntt_core.v +++ b/rtl/ip/mlkem/ntt/ntt_core.v @@ -53,8 +53,8 @@ module ntt_core ( reg rd_fire_r; // issue_fire delayed 1 (BRAM read latency) -> feed // 9-deep writeback address delay lines (BRAM read + 8-stage butterfly). - reg [7:0] a0d[0:8]; reg [7:0] b0d[0:8]; - reg [7:0] a1d[0:8]; reg [7:0] b1d[0:8]; + reg [7:0] a0d[0:9]; reg [7:0] b0d[0:9]; + reg [7:0] a1d[0:9]; reg [7:0] b1d[0:9]; integer k; reg [DW-1:0] coeff_out_r; @@ -128,10 +128,10 @@ module ntt_core ( wire wr_en_1 = bf0_valid & ~loading_we; wire wr_en_2 = bf0_valid & ~loading_we; wire wr_en_3 = bf0_valid & ~loading_we; - wire [7:0] wr_addr_0 = loading_we ? load_addr : a0d[8]; - wire [7:0] wr_addr_1 = b0d[8]; - wire [7:0] wr_addr_2 = a1d[8]; - wire [7:0] wr_addr_3 = b1d[8]; + wire [7:0] wr_addr_0 = loading_we ? load_addr : a0d[9]; + wire [7:0] wr_addr_1 = b0d[9]; + wire [7:0] wr_addr_2 = a1d[9]; + wire [7:0] wr_addr_3 = b1d[9]; wire [11:0] wr_data_0 = loading_we ? coeff_in : bf0_a_out; wire [11:0] wr_data_1 = bf0_b_out; wire [11:0] wr_data_2 = bf1_a_out; @@ -185,7 +185,7 @@ module ntt_core ( layer <= 3'd0; mode_r <= 1'b0; layer_issue_done <= 1'b0; bf_inflight <= 6'd0; rd_fire_r <= 1'b0; - for (k = 0; k <= 8; k = k + 1) begin + for (k = 0; k <= 9; k = k + 1) begin a0d[k] <= 8'd0; b0d[k] <= 8'd0; a1d[k] <= 8'd0; b1d[k] <= 8'd0; end coeff_out_r <= 12'd0; valid_o_r <= 1'b0; @@ -200,7 +200,7 @@ module ntt_core ( // writeback address delay lines: sample j-derived addrs at read cycle a0d[0] <= cmp_a0; b0d[0] <= cmp_b0; a1d[0] <= cmp_a1; b1d[0] <= cmp_b1; - for (k = 1; k <= 8; k = k + 1) begin + for (k = 1; k <= 9; k = k + 1) begin a0d[k] <= a0d[k-1]; b0d[k] <= b0d[k-1]; a1d[k] <= a1d[k-1]; b1d[k] <= b1d[k-1]; end