wu-arch/vortex
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: wu-arch:wu-blackwell-fp8-e4m3
Branches Tags
wu-arch:wu-architecture wu-arch:wu-blackwell-fp8-e4m3 wu-arch:wu-tmem-bank-scalar-softmax wu-arch:wu-blackwell-4lane-tc wu-arch:master
..
compare: wu-arch:2bfc6c4bdef3a8b33bb8b53fae5b82e44df480f0
Branches Tags
wu-arch:wu-blackwell-fp8-e4m3 wu-arch:wu-tmem-bank-scalar-softmax wu-arch:wu-blackwell-4lane-tc wu-arch:wu-architecture wu-arch:master

1 Commits
wu-blackwe ... 2bfc6c4bde

Author SHA1 Message Date
Zhongdi LUO
2bfc6c4bde fix: pack scalar tmem softmax as fp16 2026-07-10 13:01:58 +00:00
2 changed files with 127 additions and 443 deletions
Expand all files Collapse all files

183
hw/rtl/core/VX_execute.sv
View File

@@ -356,12 +356,15 @@ module VX_execute import VX_gpu_pkg::*; #(
reg scalar_tmem_store_pending; reg scalar_tmem_store_pending;
reg scalar_tmem_commit_pending; reg scalar_tmem_commit_pending;
reg scalar_tmem_softmax_active; reg scalar_tmem_softmax_active;
reg [5:0] scalar_tmem_softmax_read_index; reg scalar_tmem_softmax_read_pending;
reg scalar_tmem_softmax_reads_issued; reg [1:0] scalar_tmem_softmax_stage;
reg scalar_tmem_softmax_read_response_valid; reg [5:0] scalar_tmem_softmax_index;
reg [5:0] scalar_tmem_softmax_read_response_index; reg [5:0] scalar_tmem_softmax_read_index_r;
reg [8:0] scalar_tmem_softmax_score_base_r; reg [8:0] scalar_tmem_softmax_score_base_r;
reg [8:0] scalar_tmem_softmax_p_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 [`UUID_WIDTH-1:0] scalar_tmem_uuid_r;
reg [`NW_WIDTH-1:0] scalar_tmem_wid_r; reg [`NW_WIDTH-1:0] scalar_tmem_wid_r;
reg [`NUM_THREADS-1:0] scalar_tmem_tmask_r; reg [`NUM_THREADS-1:0] scalar_tmem_tmask_r;
@@ -378,9 +381,50 @@ 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][`NR_BITS-1:0] scalar_tmem_req_rd;
wire [`ISSUE_WIDTH-1:0][8:0] scalar_tmem_req_softmax_p_addr; 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_ROW_LAST = 6'd31;
localparam [5:0] SCALAR_TMEM_SOFTMAX_TILE_LAST = 6'd63; 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; assign scalar_tmem_pending = |scalar_tmem_dispatch;
always @(*) begin always @(*) begin
@@ -403,48 +447,22 @@ module VX_execute import VX_gpu_pkg::*; #(
scalar_tmem_grant_is_load ? scalar_tmem_rready : scalar_tmem_wready; 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_req_fire = scalar_tmem_grant_valid && scalar_tmem_req_ready;
wire scalar_tmem_softmax_read_valid = scalar_tmem_softmax_active wire scalar_tmem_softmax_read_valid = scalar_tmem_softmax_active
&& !scalar_tmem_softmax_reads_issued; && !scalar_tmem_softmax_read_pending
wire scalar_tmem_softmax_read_issue = scalar_tmem_softmax_read_valid && scalar_tmem_rready; && (scalar_tmem_softmax_stage != SCALAR_TMEM_SOFTMAX_WRITE);
wire [8:0] scalar_tmem_softmax_read_addr = scalar_tmem_softmax_score_base_r
+ 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 wire scalar_tmem_softmax_write_valid = scalar_tmem_softmax_active
&& scalar_tmem_softmax_vector_out_valid; && (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_WRITE);
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 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]);
wire [8:0] scalar_tmem_softmax_write_addr = scalar_tmem_softmax_p_base_r wire [8:0] scalar_tmem_softmax_write_addr = scalar_tmem_softmax_p_base_r
+ 9'(scalar_tmem_softmax_vector_out_index); + 9'(scalar_tmem_softmax_index);
wire [`NUM_THREADS*`XLEN-1:0] scalar_tmem_softmax_wdata = 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 wire scalar_tmem_softmax_commit = scalar_tmem_commit_pending
&& !scalar_tmem_load_pending && !scalar_tmem_load_pending
&& !scalar_tmem_store_pending; && !scalar_tmem_store_pending;
wire [`XLEN-1:0] scalar_tmem_softmax_token = `XLEN'(scalar_tmem_softmax_p_base_r); 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 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_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]; assign scalar_tmem_req_addr[i] = lsu_dispatch_if[i].data.rs1_data[0][8:0];
@@ -485,6 +503,14 @@ module VX_execute import VX_gpu_pkg::*; #(
end end
endfunction 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 `ifdef DBG_TRACE_CORE_PIPELINE_VCS
always @(posedge clk) begin always @(posedge clk) begin
if (!reset && ($time > `TRACE_STARTTIME) && (CORE_ID == 0)) begin if (!reset && ($time > `TRACE_STARTTIME) && (CORE_ID == 0)) begin
@@ -494,14 +520,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_rd[0], scalar_tmem_req_addr[0], scalar_tmem_req_softmax_p_addr[0],
scalar_tmem_req_tmask[0], scalar_tmem_req_uuid[0])); scalar_tmem_req_tmask[0], scalar_tmem_req_uuid[0]));
end end
if (scalar_tmem_softmax_read_issue && (scalar_tmem_softmax_read_index == 6'd0 || scalar_tmem_softmax_read_index == SCALAR_TMEM_SOFTMAX_ROW_LAST)) begin 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: index=%0d addr=%0d rready=%b PC=0x%0h rd=%0d\n", `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_read_index, $time, CORE_ID, scalar_tmem_softmax_stage, scalar_tmem_softmax_index,
scalar_tmem_softmax_read_addr, scalar_tmem_rready, scalar_tmem_pc_r, scalar_tmem_rd_r)); scalar_tmem_softmax_read_addr, scalar_tmem_rready, scalar_tmem_pc_r, scalar_tmem_rd_r));
end end
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 if (scalar_tmem_softmax_write_issue && (scalar_tmem_softmax_index == 6'd0 || scalar_tmem_softmax_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", `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_vector_out_index, scalar_tmem_softmax_write_addr, $time, CORE_ID, scalar_tmem_softmax_index, scalar_tmem_softmax_write_addr,
scalar_tmem_wready, scalar_tmem_pc_r, scalar_tmem_rd_r)); scalar_tmem_wready, scalar_tmem_pc_r, scalar_tmem_rd_r));
end end
if (scalar_tmem_commit_pending && scalar_tmem_pc_r == 32'h80000318) begin if (scalar_tmem_commit_pending && scalar_tmem_pc_r == 32'h80000318) begin
@@ -550,10 +576,10 @@ module VX_execute import VX_gpu_pkg::*; #(
scalar_tmem_store_pending <= 1'b0; scalar_tmem_store_pending <= 1'b0;
scalar_tmem_commit_pending <= 1'b0; scalar_tmem_commit_pending <= 1'b0;
scalar_tmem_softmax_active <= 1'b0; scalar_tmem_softmax_active <= 1'b0;
scalar_tmem_softmax_read_index <= '0; scalar_tmem_softmax_read_pending <= 1'b0;
scalar_tmem_softmax_reads_issued <= 1'b0; scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_MAX;
scalar_tmem_softmax_read_response_valid <= 1'b0; scalar_tmem_softmax_index <= '0;
scalar_tmem_softmax_read_response_index <= '0; scalar_tmem_softmax_read_index_r <= '0;
scalar_tmem_softmax_score_base_r <= '0; scalar_tmem_softmax_score_base_r <= '0;
scalar_tmem_softmax_p_base_r <= '0; scalar_tmem_softmax_p_base_r <= '0;
scalar_tmem_uuid_r <= '0; scalar_tmem_uuid_r <= '0;
@@ -579,12 +605,16 @@ module VX_execute import VX_gpu_pkg::*; #(
scalar_tmem_rd_r <= scalar_tmem_req_rd[i]; scalar_tmem_rd_r <= scalar_tmem_req_rd[i];
if (scalar_tmem_grant_is_softmax) begin if (scalar_tmem_grant_is_softmax) begin
scalar_tmem_softmax_active <= 1'b1; scalar_tmem_softmax_active <= 1'b1;
scalar_tmem_softmax_read_index <= '0; scalar_tmem_softmax_read_pending <= 1'b0;
scalar_tmem_softmax_reads_issued <= 1'b0; scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_MAX;
scalar_tmem_softmax_read_response_valid <= 1'b0; scalar_tmem_softmax_index <= '0;
scalar_tmem_softmax_read_response_index <= '0; scalar_tmem_softmax_read_index_r <= '0;
scalar_tmem_softmax_score_base_r <= scalar_tmem_req_addr[i]; scalar_tmem_softmax_score_base_r <= scalar_tmem_req_addr[i];
scalar_tmem_softmax_p_base_r <= scalar_tmem_req_softmax_p_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 end
end end
@@ -593,19 +623,56 @@ module VX_execute import VX_gpu_pkg::*; #(
scalar_tmem_rdata_valid_r <= 1'b1; scalar_tmem_rdata_valid_r <= 1'b1;
end end
scalar_tmem_softmax_read_response_valid <= scalar_tmem_softmax_read_issue;
if (scalar_tmem_softmax_read_issue) begin if (scalar_tmem_softmax_read_issue) begin
scalar_tmem_softmax_read_response_index <= scalar_tmem_softmax_read_index; scalar_tmem_softmax_read_pending <= 1'b1;
if (scalar_tmem_softmax_read_index == SCALAR_TMEM_SOFTMAX_ROW_LAST) begin scalar_tmem_softmax_read_index_r <= scalar_tmem_softmax_index;
scalar_tmem_softmax_reads_issued <= 1'b1; end else if (scalar_tmem_softmax_read_pending) begin
end else begin scalar_tmem_softmax_read_pending <= 1'b0;
scalar_tmem_softmax_read_index <= scalar_tmem_softmax_read_index + 6'd1; 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
end end
if (scalar_tmem_softmax_vector_done) begin 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
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_softmax_active <= 1'b0;
scalar_tmem_commit_pending <= 1'b1; scalar_tmem_commit_pending <= 1'b1;
end else begin
scalar_tmem_softmax_index <= scalar_tmem_softmax_index + 6'd1;
end
end end
if (scalar_tmem_commit_pending && (|(scalar_tmem_grant_r & scalar_tmem_commit_ready))) begin if (scalar_tmem_commit_pending && (|(scalar_tmem_grant_r & scalar_tmem_commit_ready))) begin

383
hw/rtl/core/VX_tmem_softmax_unit.sv
View File

@@ -1,383 +0,0 @@
// 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
// Convert a positive unsigned Q15 probability to FP8 E4M3 using
// round-to-nearest-even. Softmax probabilities are in [0, 1], so the
// positive subnormal and exponent ranges through 1.0 are sufficient.
function automatic [7:0] q15_to_fp8e4m3;
input [15:0] value;
integer msb;
integer i;
integer shift;
reg [3:0] exponent;
reg [4:0] rounded;
reg [15:0] remainder;
reg [15:0] halfway;
begin
if (value < 16'd512) begin
// E4M3 subnormals and the minimum normal are spaced by 2^-9,
// which is 64 in unsigned Q15.
rounded = value[10:6];
remainder = value & 16'h003f;
if ((remainder > 16'd32)
|| ((remainder == 16'd32) && rounded[0]))
rounded = rounded + 1'b1;
q15_to_fp8e4m3 = {3'b000, rounded};
end else begin
msb = 0;
for (i = 0; i < 16; ++i) begin
if (value[i]) msb = i;
end
exponent = 4'(msb - 8);
shift = integer'(exponent) + 5;
rounded = 5'(value >> shift);
remainder = value & ((16'd1 << shift) - 1'b1);
halfway = 16'd1 << (shift - 1);
if ((remainder > halfway)
|| ((remainder == halfway) && rounded[0]))
rounded = rounded + 1'b1;
if (rounded == 5'd16)
q15_to_fp8e4m3 = {1'b0, exponent + 1'b1, 3'b000};
else
q15_to_fp8e4m3 = {1'b0, exponent, rounded[2:0]};
end
end
endfunction
function automatic [7:0] normalize_to_fp8e4m3;
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 > 56'd32768)
normalize_to_fp8e4m3 = 8'h38;
else
normalize_to_fp8e4m3 = q15_to_fp8e4m3(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] <= {
4{normalize_to_fp8e4m3(
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