From eb1b2ea330ccb24d8e5fadcc517dce9c267c73a2 Mon Sep 17 00:00:00 2001 From: Zhongdi LUO Date: Sun, 12 Jul 2026 01:59:32 +0000 Subject: [PATCH] feat: add pipelined scalar tmem softmax unit --- hw/rtl/core/VX_execute.sv | 187 +++++--------- hw/rtl/core/VX_tmem_softmax_unit.sv | 362 ++++++++++++++++++++++++++++ 2 files changed, 422 insertions(+), 127 deletions(-) create mode 100644 hw/rtl/core/VX_tmem_softmax_unit.sv diff --git a/hw/rtl/core/VX_execute.sv b/hw/rtl/core/VX_execute.sv index 5077bd25..99d89371 100644 --- a/hw/rtl/core/VX_execute.sv +++ b/hw/rtl/core/VX_execute.sv @@ -356,15 +356,12 @@ module VX_execute import VX_gpu_pkg::*; #( reg scalar_tmem_store_pending; reg scalar_tmem_commit_pending; reg scalar_tmem_softmax_active; - reg scalar_tmem_softmax_read_pending; - reg [1:0] scalar_tmem_softmax_stage; - reg [5:0] scalar_tmem_softmax_index; - reg [5:0] scalar_tmem_softmax_read_index_r; + reg [5:0] scalar_tmem_softmax_read_index; + reg scalar_tmem_softmax_reads_issued; + reg scalar_tmem_softmax_read_response_valid; + reg [5:0] scalar_tmem_softmax_read_response_index; reg [8:0] scalar_tmem_softmax_score_base_r; reg [8:0] scalar_tmem_softmax_p_base_r; - reg [31:0] scalar_tmem_softmax_scores_r [0:31][0:`NUM_THREADS-1]; - reg [31:0] scalar_tmem_softmax_row_max_r [0:`NUM_THREADS-1]; - reg [31:0] scalar_tmem_softmax_denom_r [0:`NUM_THREADS-1]; reg [`UUID_WIDTH-1:0] scalar_tmem_uuid_r; reg [`NW_WIDTH-1:0] scalar_tmem_wid_r; reg [`NUM_THREADS-1:0] scalar_tmem_tmask_r; @@ -381,50 +378,9 @@ module VX_execute import VX_gpu_pkg::*; #( wire [`ISSUE_WIDTH-1:0][`NR_BITS-1:0] scalar_tmem_req_rd; wire [`ISSUE_WIDTH-1:0][8:0] scalar_tmem_req_softmax_p_addr; - localparam [1:0] SCALAR_TMEM_SOFTMAX_MAX = 2'd0; - localparam [1:0] SCALAR_TMEM_SOFTMAX_DENOM = 2'd1; - localparam [1:0] SCALAR_TMEM_SOFTMAX_WRITE = 2'd2; localparam [5:0] SCALAR_TMEM_SOFTMAX_ROW_LAST = 6'd31; localparam [5:0] SCALAR_TMEM_SOFTMAX_TILE_LAST = 6'd63; - function automatic [31:0] scalar_tmem_f32_max; - input [31:0] a_bits; - input [31:0] b_bits; - begin - `ifdef SV_DPI - scalar_tmem_f32_max = 32'(dpi_f32_max(1'b1, int'(a_bits), int'(b_bits))); - `else - scalar_tmem_f32_max = a_bits; - `endif - end - endfunction - - function automatic [31:0] scalar_tmem_softmax_exp_acc; - input [31:0] score_bits; - input [31:0] max_bits; - input [31:0] accum_bits; - begin - `ifdef SV_DPI - scalar_tmem_softmax_exp_acc = 32'(dpi_softmax_exp_acc(1'b1, int'(score_bits), int'(max_bits), int'(accum_bits))); - `else - scalar_tmem_softmax_exp_acc = accum_bits; - `endif - end - endfunction - - function automatic [31:0] scalar_tmem_softmax_prob_pack; - input [31:0] score_bits; - input [31:0] max_bits; - input [31:0] denom_bits; - begin - `ifdef SV_DPI - scalar_tmem_softmax_prob_pack = 32'(dpi_softmax_prob_to_f16x2(1'b1, int'(score_bits), int'(max_bits), int'(denom_bits))); - `else - scalar_tmem_softmax_prob_pack = '0; - `endif - end - endfunction - assign scalar_tmem_pending = |scalar_tmem_dispatch; always @(*) begin @@ -447,22 +403,48 @@ module VX_execute import VX_gpu_pkg::*; #( scalar_tmem_grant_is_load ? scalar_tmem_rready : scalar_tmem_wready; wire scalar_tmem_req_fire = scalar_tmem_grant_valid && scalar_tmem_req_ready; wire scalar_tmem_softmax_read_valid = scalar_tmem_softmax_active - && !scalar_tmem_softmax_read_pending - && (scalar_tmem_softmax_stage != SCALAR_TMEM_SOFTMAX_WRITE); - wire scalar_tmem_softmax_write_valid = scalar_tmem_softmax_active - && (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_WRITE); + && !scalar_tmem_softmax_reads_issued; wire scalar_tmem_softmax_read_issue = scalar_tmem_softmax_read_valid && scalar_tmem_rready; - wire scalar_tmem_softmax_write_issue = scalar_tmem_softmax_write_valid && scalar_tmem_wready; wire [8:0] scalar_tmem_softmax_read_addr = scalar_tmem_softmax_score_base_r - + 9'(scalar_tmem_softmax_index[4:0]); + + 9'(scalar_tmem_softmax_read_index[4:0]); + + wire scalar_tmem_softmax_vector_busy; + wire scalar_tmem_softmax_vector_out_valid; + wire [5:0] scalar_tmem_softmax_vector_out_index; + wire [`NUM_THREADS*`XLEN-1:0] scalar_tmem_softmax_vector_out_data; + wire scalar_tmem_softmax_vector_done; + wire scalar_tmem_softmax_vector_out_ready = scalar_tmem_softmax_active && scalar_tmem_wready; + wire scalar_tmem_softmax_write_valid = scalar_tmem_softmax_active + && scalar_tmem_softmax_vector_out_valid; + wire scalar_tmem_softmax_write_issue = scalar_tmem_softmax_write_valid && scalar_tmem_wready; wire [8:0] scalar_tmem_softmax_write_addr = scalar_tmem_softmax_p_base_r - + 9'(scalar_tmem_softmax_index); - wire [`NUM_THREADS*`XLEN-1:0] scalar_tmem_softmax_wdata; + + 9'(scalar_tmem_softmax_vector_out_index); + wire [`NUM_THREADS*`XLEN-1:0] scalar_tmem_softmax_wdata = + scalar_tmem_softmax_vector_out_data; wire scalar_tmem_softmax_commit = scalar_tmem_commit_pending && !scalar_tmem_load_pending && !scalar_tmem_store_pending; wire [`XLEN-1:0] scalar_tmem_softmax_token = `XLEN'(scalar_tmem_softmax_p_base_r); + VX_tmem_softmax_unit #( + .NUM_LANES (`NUM_THREADS), + .ROW_SIZE (32) + ) tmem_softmax_unit ( + .clk (clk), + .reset (reset), + .start (scalar_tmem_req_fire && scalar_tmem_grant_is_softmax), + .load_valid (scalar_tmem_softmax_read_response_valid), + .load_data (scalar_tmem_rdata), + .load_last (scalar_tmem_softmax_read_response_index == SCALAR_TMEM_SOFTMAX_ROW_LAST), + .busy (scalar_tmem_softmax_vector_busy), + .out_valid (scalar_tmem_softmax_vector_out_valid), + .out_ready (scalar_tmem_softmax_vector_out_ready), + .out_index (scalar_tmem_softmax_vector_out_index), + .out_data (scalar_tmem_softmax_vector_out_data), + .done (scalar_tmem_softmax_vector_done) + ); + `UNUSED_VAR (scalar_tmem_softmax_vector_busy) + for (genvar i = 0; i < `ISSUE_WIDTH; ++i) begin : g_scalar_tmem_ready assign scalar_tmem_dispatch_ready[i] = scalar_tmem_grant_valid && scalar_tmem_grant[i] && scalar_tmem_req_ready; assign scalar_tmem_req_addr[i] = lsu_dispatch_if[i].data.rs1_data[0][8:0]; @@ -503,14 +485,6 @@ module VX_execute import VX_gpu_pkg::*; #( end endfunction - for (genvar lane = 0; lane < `NUM_THREADS; ++lane) begin : g_scalar_tmem_softmax_wdata - assign scalar_tmem_softmax_wdata[lane * `XLEN +: `XLEN] = - scalar_tmem_softmax_prob_pack( - scalar_tmem_softmax_scores_r[scalar_tmem_softmax_index[4:0]][lane], - scalar_tmem_softmax_row_max_r[lane], - scalar_tmem_softmax_denom_r[lane]); - end - `ifdef DBG_TRACE_CORE_PIPELINE_VCS always @(posedge clk) begin if (!reset && ($time > `TRACE_STARTTIME) && (CORE_ID == 0)) begin @@ -520,14 +494,14 @@ module VX_execute import VX_gpu_pkg::*; #( scalar_tmem_req_rd[0], scalar_tmem_req_addr[0], scalar_tmem_req_softmax_p_addr[0], scalar_tmem_req_tmask[0], scalar_tmem_req_uuid[0])); end - if (scalar_tmem_softmax_read_issue && (scalar_tmem_softmax_index == 6'd0 || scalar_tmem_softmax_index == SCALAR_TMEM_SOFTMAX_ROW_LAST)) begin - `TRACE(1, ("%d: core%0d-scalar-tmem-softmax-read: stage=%0d index=%0d addr=%0d rready=%b PC=0x%0h rd=%0d\n", - $time, CORE_ID, scalar_tmem_softmax_stage, scalar_tmem_softmax_index, + if (scalar_tmem_softmax_read_issue && (scalar_tmem_softmax_read_index == 6'd0 || scalar_tmem_softmax_read_index == SCALAR_TMEM_SOFTMAX_ROW_LAST)) begin + `TRACE(1, ("%d: core%0d-scalar-tmem-softmax-read: index=%0d addr=%0d rready=%b PC=0x%0h rd=%0d\n", + $time, CORE_ID, scalar_tmem_softmax_read_index, scalar_tmem_softmax_read_addr, scalar_tmem_rready, scalar_tmem_pc_r, scalar_tmem_rd_r)); end - if (scalar_tmem_softmax_write_issue && (scalar_tmem_softmax_index == 6'd0 || scalar_tmem_softmax_index == SCALAR_TMEM_SOFTMAX_TILE_LAST)) begin + if (scalar_tmem_softmax_write_issue && (scalar_tmem_softmax_vector_out_index == 6'd0 || scalar_tmem_softmax_vector_out_index == SCALAR_TMEM_SOFTMAX_TILE_LAST)) begin `TRACE(1, ("%d: core%0d-scalar-tmem-softmax-write: index=%0d addr=%0d wready=%b PC=0x%0h rd=%0d\n", - $time, CORE_ID, scalar_tmem_softmax_index, scalar_tmem_softmax_write_addr, + $time, CORE_ID, scalar_tmem_softmax_vector_out_index, scalar_tmem_softmax_write_addr, scalar_tmem_wready, scalar_tmem_pc_r, scalar_tmem_rd_r)); end if (scalar_tmem_commit_pending && scalar_tmem_pc_r == 32'h80000318) begin @@ -576,10 +550,10 @@ module VX_execute import VX_gpu_pkg::*; #( scalar_tmem_store_pending <= 1'b0; scalar_tmem_commit_pending <= 1'b0; scalar_tmem_softmax_active <= 1'b0; - scalar_tmem_softmax_read_pending <= 1'b0; - scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_MAX; - scalar_tmem_softmax_index <= '0; - scalar_tmem_softmax_read_index_r <= '0; + scalar_tmem_softmax_read_index <= '0; + scalar_tmem_softmax_reads_issued <= 1'b0; + scalar_tmem_softmax_read_response_valid <= 1'b0; + scalar_tmem_softmax_read_response_index <= '0; scalar_tmem_softmax_score_base_r <= '0; scalar_tmem_softmax_p_base_r <= '0; scalar_tmem_uuid_r <= '0; @@ -605,16 +579,12 @@ module VX_execute import VX_gpu_pkg::*; #( scalar_tmem_rd_r <= scalar_tmem_req_rd[i]; if (scalar_tmem_grant_is_softmax) begin scalar_tmem_softmax_active <= 1'b1; - scalar_tmem_softmax_read_pending <= 1'b0; - scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_MAX; - scalar_tmem_softmax_index <= '0; - scalar_tmem_softmax_read_index_r <= '0; + scalar_tmem_softmax_read_index <= '0; + scalar_tmem_softmax_reads_issued <= 1'b0; + scalar_tmem_softmax_read_response_valid <= 1'b0; + scalar_tmem_softmax_read_response_index <= '0; scalar_tmem_softmax_score_base_r <= scalar_tmem_req_addr[i]; scalar_tmem_softmax_p_base_r <= scalar_tmem_req_softmax_p_addr[i]; - for (integer lane = 0; lane < `NUM_THREADS; ++lane) begin - scalar_tmem_softmax_row_max_r[lane] <= '0; - scalar_tmem_softmax_denom_r[lane] <= '0; - end end end end @@ -623,56 +593,19 @@ module VX_execute import VX_gpu_pkg::*; #( scalar_tmem_rdata_valid_r <= 1'b1; end + scalar_tmem_softmax_read_response_valid <= scalar_tmem_softmax_read_issue; if (scalar_tmem_softmax_read_issue) begin - scalar_tmem_softmax_read_pending <= 1'b1; - scalar_tmem_softmax_read_index_r <= scalar_tmem_softmax_index; - end else if (scalar_tmem_softmax_read_pending) begin - scalar_tmem_softmax_read_pending <= 1'b0; - for (integer lane = 0; lane < `NUM_THREADS; ++lane) begin - scalar_tmem_softmax_scores_r[scalar_tmem_softmax_read_index_r[4:0]][lane] - <= scalar_tmem_rdata[lane * `XLEN +: `XLEN]; - if (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_MAX) begin - scalar_tmem_softmax_row_max_r[lane] <= - (scalar_tmem_softmax_read_index_r == '0) - ? scalar_tmem_rdata[lane * `XLEN +: `XLEN] - : scalar_tmem_f32_max(scalar_tmem_softmax_row_max_r[lane], - scalar_tmem_rdata[lane * `XLEN +: `XLEN]); - end else if (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_DENOM) begin - scalar_tmem_softmax_denom_r[lane] <= - scalar_tmem_softmax_exp_acc( - scalar_tmem_rdata[lane * `XLEN +: `XLEN], - scalar_tmem_softmax_row_max_r[lane], - scalar_tmem_softmax_denom_r[lane]); - end - end - - if (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_MAX) begin - if (scalar_tmem_softmax_read_index_r == SCALAR_TMEM_SOFTMAX_ROW_LAST) begin - scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_DENOM; - scalar_tmem_softmax_index <= '0; - for (integer lane = 0; lane < `NUM_THREADS; ++lane) begin - scalar_tmem_softmax_denom_r[lane] <= 32'h00000000; - end - end else begin - scalar_tmem_softmax_index <= scalar_tmem_softmax_read_index_r + 6'd1; - end - end else if (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_DENOM) begin - if (scalar_tmem_softmax_read_index_r == SCALAR_TMEM_SOFTMAX_ROW_LAST) begin - scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_WRITE; - scalar_tmem_softmax_index <= '0; - end else begin - scalar_tmem_softmax_index <= scalar_tmem_softmax_read_index_r + 6'd1; - end + scalar_tmem_softmax_read_response_index <= scalar_tmem_softmax_read_index; + if (scalar_tmem_softmax_read_index == SCALAR_TMEM_SOFTMAX_ROW_LAST) begin + scalar_tmem_softmax_reads_issued <= 1'b1; + end else begin + scalar_tmem_softmax_read_index <= scalar_tmem_softmax_read_index + 6'd1; end end - if (scalar_tmem_softmax_write_issue) begin - if (scalar_tmem_softmax_index == SCALAR_TMEM_SOFTMAX_TILE_LAST) begin - scalar_tmem_softmax_active <= 1'b0; - scalar_tmem_commit_pending <= 1'b1; - end else begin - scalar_tmem_softmax_index <= scalar_tmem_softmax_index + 6'd1; - end + if (scalar_tmem_softmax_vector_done) begin + scalar_tmem_softmax_active <= 1'b0; + scalar_tmem_commit_pending <= 1'b1; end if (scalar_tmem_commit_pending && (|(scalar_tmem_grant_r & scalar_tmem_commit_ready))) begin diff --git a/hw/rtl/core/VX_tmem_softmax_unit.sv b/hw/rtl/core/VX_tmem_softmax_unit.sv new file mode 100644 index 00000000..fd38f929 --- /dev/null +++ b/hw/rtl/core/VX_tmem_softmax_unit.sv @@ -0,0 +1,362 @@ +// Streaming, synthesizable softmax engine for the Blackwell TMEM path. +// Four independent lane streams are processed in parallel; each stream has +// ROW_SIZE elements. The exponential datapath is pipelined with II=1. +module VX_tmem_softmax_unit #( + parameter integer NUM_LANES = 4, + parameter integer ROW_SIZE = 32 +) ( + input wire clk, + input wire reset, + input wire start, + input wire load_valid, + input wire [NUM_LANES-1:0][31:0] load_data, + input wire load_last, + output wire busy, + output reg out_valid, + input wire out_ready, + output reg [5:0] out_index, + output reg [NUM_LANES-1:0][31:0] out_data, + output reg done +); + localparam [2:0] STATE_IDLE = 3'd0; + localparam [2:0] STATE_LOAD = 3'd1; + localparam [2:0] STATE_EXP = 3'd2; + localparam [2:0] STATE_RECIP = 3'd3; + localparam [2:0] STATE_NORM = 3'd4; + + reg [2:0] state; + reg [$clog2(ROW_SIZE)-1:0] load_index; + reg [$clog2(ROW_SIZE)-1:0] exp_feed_index; + reg exp_feed_done; + reg [1:0] reciprocal_lane; + reg [5:0] norm_index; + + reg [31:0] scores [0:ROW_SIZE-1][0:NUM_LANES-1]; + reg [31:0] max_bits [0:NUM_LANES-1]; + reg [23:0] exp_values [0:ROW_SIZE-1][0:NUM_LANES-1]; + reg [29:0] exp_sums [0:NUM_LANES-1]; + reg [31:0] reciprocals [0:NUM_LANES-1]; + + reg [2:0] exp_pipe_valid; + reg [$clog2(ROW_SIZE)-1:0] exp_pipe_index [0:2]; + reg [3:0] exp_int_stage0 [0:NUM_LANES-1]; + reg [3:0] exp_frac_stage0 [0:NUM_LANES-1]; + reg [7:0] exp_rem_stage0 [0:NUM_LANES-1]; + reg [3:0] exp_int_stage1 [0:NUM_LANES-1]; + reg [23:0] exp_frac_stage1 [0:NUM_LANES-1]; + reg [23:0] exp_stage2 [0:NUM_LANES-1]; + + assign busy = state != STATE_IDLE; + + function automatic f32_greater; + input [31:0] a; + input [31:0] b; + begin + if ((a[30:0] == 0) && (b[30:0] == 0)) begin + f32_greater = 1'b0; + end else if (a[31] != b[31]) begin + f32_greater = !a[31]; + end else if (!a[31]) begin + f32_greater = a[30:0] > b[30:0]; + end else begin + f32_greater = a[30:0] < b[30:0]; + end + end + endfunction + + // Convert finite FP32 to signed Q12, saturating outside the supported + // softmax score range. Subnormals are treated as zero. + function automatic signed [23:0] f32_to_q12; + input [31:0] value; + reg [7:0] exponent; + reg [23:0] mantissa; + reg [63:0] magnitude; + integer shift; + begin + exponent = value[30:23]; + mantissa = {1'b1, value[22:0]}; + magnitude = 0; + if (exponent == 0) begin + f32_to_q12 = 24'sd0; + end else if (exponent == 8'hff) begin + f32_to_q12 = value[31] ? -24'sh7fffff : 24'sh7fffff; + end else begin + shift = integer'(exponent) - 138; + if (shift >= 0) begin + magnitude = {40'd0, mantissa} << shift; + end else if (shift > -64) begin + magnitude = {40'd0, mantissa} >> (-shift); + end + if (magnitude > 64'h7fffff) + magnitude = 64'h7fffff; + f32_to_q12 = value[31] + ? -$signed({1'b0, magnitude[22:0]}) + : $signed({1'b0, magnitude[22:0]}); + end + end + endfunction + + // {integer_part, fractional_lut_index, interpolation_remainder} for + // exp(score - max). Inputs below -8 are clamped to exp(-8). + function automatic [15:0] exp_range_info; + input signed [23:0] score_q12; + input signed [23:0] max_q12; + reg signed [24:0] distance; + reg [15:0] clamped; + begin + distance = max_q12 - score_q12; + if (distance <= 0) + clamped = 16'd0; + else if (distance >= 25'sd32768) + clamped = 16'd32768; + else + clamped = distance[15:0]; + + if (clamped == 16'd32768) + exp_range_info = {4'd8, 4'd0, 8'd0}; + else + exp_range_info = {clamped[15:12], clamped[11:8], clamped[7:0]}; + end + endfunction + + function automatic [23:0] exp_frac_lut; + input [4:0] index; + begin + case (index) + 5'd0: exp_frac_lut = 24'd1048576; + 5'd1: exp_frac_lut = 24'd985046; + 5'd2: exp_frac_lut = 24'd925365; + 5'd3: exp_frac_lut = 24'd869300; + 5'd4: exp_frac_lut = 24'd816632; + 5'd5: exp_frac_lut = 24'd767155; + 5'd6: exp_frac_lut = 24'd720675; + 5'd7: exp_frac_lut = 24'd677012; + 5'd8: exp_frac_lut = 24'd635993; + 5'd9: exp_frac_lut = 24'd597461; + 5'd10: exp_frac_lut = 24'd561262; + 5'd11: exp_frac_lut = 24'd527257; + 5'd12: exp_frac_lut = 24'd495312; + 5'd13: exp_frac_lut = 24'd465303; + 5'd14: exp_frac_lut = 24'd437112; + 5'd15: exp_frac_lut = 24'd410628; + default: exp_frac_lut = 24'd385750; + endcase + end + endfunction + + function automatic [23:0] exp_int_lut; + input [3:0] index; + begin + case (index) + 4'd0: exp_int_lut = 24'd1048576; + 4'd1: exp_int_lut = 24'd385750; + 4'd2: exp_int_lut = 24'd141909; + 4'd3: exp_int_lut = 24'd52206; + 4'd4: exp_int_lut = 24'd19205; + 4'd5: exp_int_lut = 24'd7065; + 4'd6: exp_int_lut = 24'd2599; + 4'd7: exp_int_lut = 24'd956; + default: exp_int_lut = 24'd352; + endcase + end + endfunction + + function automatic [15:0] q15_to_fp16; + input [15:0] value; + integer msb; + integer i; + reg [15:0] normalized; + begin + msb = -1; + for (i = 0; i < 16; ++i) begin + if (value[i]) msb = i; + end + if (msb < 0) begin + q15_to_fp16 = 16'd0; + end else begin + if (msb >= 10) + normalized = value >> (msb - 10); + else + normalized = value << (10 - msb); + q15_to_fp16 = {1'b0, msb[4:0], normalized[9:0]}; + end + end + endfunction + + function automatic [15:0] normalize_to_fp16; + input [23:0] exp_value; + input [31:0] reciprocal; + reg [55:0] product; + reg [55:0] probability; + begin + product = exp_value * reciprocal; + probability = (product + 56'd524288) >> 20; + if (probability > 36'd32768) + normalize_to_fp16 = 16'h3c00; + else + normalize_to_fp16 = q15_to_fp16(probability[15:0]); + end + endfunction + + function automatic [23:0] interpolate_exp_frac; + input [3:0] frac_index; + input [7:0] remainder; + reg [23:0] upper; + reg [23:0] lower; + reg [31:0] product; + begin + upper = exp_frac_lut({1'b0, frac_index}); + lower = exp_frac_lut({1'b0, frac_index} + 5'd1); + product = (upper - lower) * remainder; + interpolate_exp_frac = upper - ((product + 32'd128) >> 8); + end + endfunction + + function automatic [23:0] scale_exp_frac; + input [23:0] frac_value; + input [3:0] int_index; + reg [47:0] product; + begin + product = frac_value * exp_int_lut(int_index); + scale_exp_frac = (product + 48'd524288) >> 20; + end + endfunction + + function automatic [31:0] reciprocal_q15; + input [29:0] sum_value; + reg [63:0] numerator; + reg [63:0] quotient; + begin + numerator = (64'd1 << 35) + ({34'd0, sum_value} >> 1); + quotient = (sum_value == 0) ? 64'd0 : numerator / {34'd0, sum_value}; + reciprocal_q15 = quotient[31:0]; + end + endfunction + + wire exp_feed = (state == STATE_EXP) && !exp_feed_done; + wire [15:0] exp_feed_range [0:NUM_LANES-1]; + for (genvar lane = 0; lane < NUM_LANES; ++lane) begin : g_exp_feed_range + assign exp_feed_range[lane] = exp_range_info( + f32_to_q12(scores[exp_feed_index][lane]), + f32_to_q12(max_bits[lane])); + end + + integer lane; + + always @(posedge clk) begin + if (reset) begin + state <= STATE_IDLE; + load_index <= '0; + exp_feed_index <= '0; + exp_feed_done <= 1'b0; + reciprocal_lane <= '0; + norm_index <= '0; + exp_pipe_valid <= '0; + out_valid <= 1'b0; + out_index <= '0; + out_data <= '0; + done <= 1'b0; + end else begin + done <= 1'b0; + + if ((state == STATE_IDLE) && start) begin + state <= STATE_LOAD; + load_index <= '0; + exp_pipe_valid <= '0; + out_valid <= 1'b0; + end + + if ((state == STATE_LOAD) && load_valid) begin + for (lane = 0; lane < NUM_LANES; ++lane) begin + scores[load_index][lane] <= load_data[lane]; + if ((load_index == 0) || f32_greater(load_data[lane], max_bits[lane])) + max_bits[lane] <= load_data[lane]; + end + if (load_last) begin + state <= STATE_EXP; + exp_feed_index <= '0; + exp_feed_done <= 1'b0; + exp_pipe_valid <= '0; + for (lane = 0; lane < NUM_LANES; ++lane) + exp_sums[lane] <= '0; + end else begin + load_index <= load_index + 1'b1; + end + end + + exp_pipe_valid[0] <= exp_feed; + exp_pipe_valid[1] <= exp_pipe_valid[0]; + exp_pipe_valid[2] <= exp_pipe_valid[1]; + if (exp_feed) begin + exp_pipe_index[0] <= exp_feed_index; + for (lane = 0; lane < NUM_LANES; ++lane) begin + exp_int_stage0[lane] <= exp_feed_range[lane][15:12]; + exp_frac_stage0[lane] <= exp_feed_range[lane][11:8]; + exp_rem_stage0[lane] <= exp_feed_range[lane][7:0]; + end + if (exp_feed_index == $clog2(ROW_SIZE)'(ROW_SIZE - 1)) + exp_feed_done <= 1'b1; + else + exp_feed_index <= exp_feed_index + 1'b1; + end + + if (exp_pipe_valid[0]) begin + exp_pipe_index[1] <= exp_pipe_index[0]; + for (lane = 0; lane < NUM_LANES; ++lane) begin + exp_frac_stage1[lane] <= interpolate_exp_frac( + exp_frac_stage0[lane], exp_rem_stage0[lane]); + exp_int_stage1[lane] <= exp_int_stage0[lane]; + end + end + + if (exp_pipe_valid[1]) begin + exp_pipe_index[2] <= exp_pipe_index[1]; + for (lane = 0; lane < NUM_LANES; ++lane) begin + exp_stage2[lane] <= scale_exp_frac( + exp_frac_stage1[lane], exp_int_stage1[lane]); + end + end + + if (exp_pipe_valid[2]) begin + for (lane = 0; lane < NUM_LANES; ++lane) begin + exp_values[exp_pipe_index[2]][lane] <= exp_stage2[lane]; + exp_sums[lane] <= exp_sums[lane] + exp_stage2[lane]; + end + if (exp_pipe_index[2] == $clog2(ROW_SIZE)'(ROW_SIZE - 1)) begin + state <= STATE_RECIP; + reciprocal_lane <= '0; + end + end + + if (state == STATE_RECIP) begin + reciprocals[reciprocal_lane] <= reciprocal_q15(exp_sums[reciprocal_lane]); + if (reciprocal_lane == 2'(NUM_LANES - 1)) begin + state <= STATE_NORM; + norm_index <= '0; + out_valid <= 1'b0; + end else begin + reciprocal_lane <= reciprocal_lane + 1'b1; + end + end + + if (state == STATE_NORM) begin + if (out_valid && out_ready && (out_index == 6'd63)) begin + out_valid <= 1'b0; + done <= 1'b1; + state <= STATE_IDLE; + end else if (!out_valid || out_ready) begin + for (lane = 0; lane < NUM_LANES; ++lane) begin + out_data[lane] <= { + 2{normalize_to_fp16( + exp_values[norm_index[4:0]][lane], reciprocals[lane])} + }; + end + out_index <= norm_index; + out_valid <= 1'b1; + if (norm_index != 6'd63) + norm_index <= norm_index + 1'b1; + end + end + end + end +endmodule