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

feat: add pipelined fp8 scalar tmem softmax unit

Browse Source
This commit is contained in:
Zhongdi LUO
2026-07-12 01:59:32 +00:00
parent dff1107bf5
commit 331c4decaa
2 changed files with 443 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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_fp8e4m3x4(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