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

feat: add scalar tmem softmax pipeline

Browse Source
This commit is contained in:
Zhongdi LUO
2026-07-10 08:23:22 +00:00
parent 9251ba0a24
commit dff1107bf5
5 changed files with 302 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
hw/rtl/core/VX_decode.sv
View File

@@ -550,6 +550,16 @@ module VX_decode #(
end
end
`endif
3'h3: begin
if (func7 == 7'h30) begin
ex_type = `EX_LSU;
op_type = `INST_LSU_TMEM_SOFTMAX;
use_rd = 1;
`USED_IREG (rd);
`USED_IREG (rs1);
`USED_IREG (rs2);
end
end
default:;
endcase
end

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

@@ -12,6 +12,9 @@
// limitations under the License.
`include "VX_define.vh"
`ifdef SV_DPI
`include "float_dpi.vh"
`endif
module VX_execute import VX_gpu_pkg::*; #(
parameter CORE_ID = 0,
@@ -304,6 +307,7 @@ module VX_execute import VX_gpu_pkg::*; #(
wire [`ISSUE_WIDTH-1:0] scalar_tmem_dispatch;
wire [`ISSUE_WIDTH-1:0] scalar_tmem_ld_dispatch;
wire [`ISSUE_WIDTH-1:0] scalar_tmem_st_dispatch;
wire [`ISSUE_WIDTH-1:0] scalar_tmem_softmax_dispatch;
wire [`ISSUE_WIDTH-1:0] scalar_tmem_dispatch_ready;
wire [`ISSUE_WIDTH-1:0] scalar_tmem_dispatch_fire;
wire [`ISSUE_WIDTH-1:0] scalar_tmem_commit_ready;
@@ -314,7 +318,11 @@ module VX_execute import VX_gpu_pkg::*; #(
&& (lsu_dispatch_if[i].data.op_type == `INST_LSU_TMEM_LD);
assign scalar_tmem_st_dispatch[i] = lsu_dispatch_if[i].valid
&& (lsu_dispatch_if[i].data.op_type == `INST_LSU_TMEM_ST);
assign scalar_tmem_dispatch[i] = scalar_tmem_ld_dispatch[i] || scalar_tmem_st_dispatch[i];
assign scalar_tmem_softmax_dispatch[i] = lsu_dispatch_if[i].valid
&& (lsu_dispatch_if[i].data.op_type == `INST_LSU_TMEM_SOFTMAX);
assign scalar_tmem_dispatch[i] = scalar_tmem_ld_dispatch[i]
|| scalar_tmem_st_dispatch[i]
|| scalar_tmem_softmax_dispatch[i];
assign scalar_mem_lsu_dispatch_if[i].valid = lsu_dispatch_if[i].valid && !scalar_tmem_dispatch[i];
assign scalar_mem_lsu_dispatch_if[i].data = lsu_dispatch_if[i].data;
@@ -347,6 +355,16 @@ module VX_execute import VX_gpu_pkg::*; #(
reg scalar_tmem_load_pending;
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 [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;
@@ -361,6 +379,51 @@ module VX_execute import VX_gpu_pkg::*; #(
wire [`ISSUE_WIDTH-1:0][`NUM_THREADS-1:0] scalar_tmem_req_tmask;
wire [`ISSUE_WIDTH-1:0][`XLEN-1:0] scalar_tmem_req_pc;
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;
@@ -375,11 +438,30 @@ module VX_execute import VX_gpu_pkg::*; #(
end
wire scalar_tmem_grant_valid = scalar_tmem_pending && !scalar_tmem_load_pending
&& !scalar_tmem_store_pending && !scalar_tmem_commit_pending;
&& !scalar_tmem_store_pending && !scalar_tmem_commit_pending
&& !scalar_tmem_softmax_active;
wire scalar_tmem_grant_is_load = |(scalar_tmem_grant & scalar_tmem_ld_dispatch);
wire scalar_tmem_grant_is_store = |(scalar_tmem_grant & scalar_tmem_st_dispatch);
wire scalar_tmem_req_ready = scalar_tmem_grant_is_load ? scalar_tmem_rready : scalar_tmem_wready;
wire scalar_tmem_grant_is_softmax = |(scalar_tmem_grant & scalar_tmem_softmax_dispatch);
wire scalar_tmem_req_ready = scalar_tmem_grant_is_softmax ? 1'b1 :
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);
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]);
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;
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);
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;
@@ -390,15 +472,19 @@ module VX_execute import VX_gpu_pkg::*; #(
assign scalar_tmem_req_tmask[i] = lsu_dispatch_if[i].data.tmask;
assign scalar_tmem_req_pc[i] = lsu_dispatch_if[i].data.PC;
assign scalar_tmem_req_rd[i] = lsu_dispatch_if[i].data.rd;
assign scalar_tmem_req_softmax_p_addr[i] = lsu_dispatch_if[i].data.rs2_data[0][8:0];
assign scalar_tmem_commit_if[i].valid = scalar_tmem_commit_pending && scalar_tmem_grant_r[i];
assign scalar_tmem_commit_if[i].data.uuid = scalar_tmem_uuid_r;
assign scalar_tmem_commit_if[i].data.wid = scalar_tmem_wid_r;
assign scalar_tmem_commit_if[i].data.tmask = scalar_tmem_tmask_r;
assign scalar_tmem_commit_if[i].data.PC = scalar_tmem_pc_r;
assign scalar_tmem_commit_if[i].data.wb = scalar_tmem_load_pending;
assign scalar_tmem_commit_if[i].data.rd = scalar_tmem_load_pending ? scalar_tmem_rd_r : '0;
assign scalar_tmem_commit_if[i].data.data = scalar_tmem_rdata_valid_r ? scalar_tmem_rdata_r : scalar_tmem_rdata;
assign scalar_tmem_commit_if[i].data.wb = scalar_tmem_load_pending || scalar_tmem_softmax_commit;
assign scalar_tmem_commit_if[i].data.rd = (scalar_tmem_load_pending || scalar_tmem_softmax_commit) ? scalar_tmem_rd_r : '0;
assign scalar_tmem_commit_if[i].data.data =
scalar_tmem_load_pending
? (scalar_tmem_rdata_valid_r ? scalar_tmem_rdata_r : scalar_tmem_rdata)
: {`NUM_THREADS{scalar_tmem_softmax_token}};
assign scalar_tmem_commit_if[i].data.tensor = 1'b0;
assign scalar_tmem_commit_if[i].data.pid = '0;
assign scalar_tmem_commit_if[i].data.sop = 1'b1;
@@ -417,9 +503,49 @@ 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
if (scalar_tmem_req_fire && scalar_tmem_grant_is_softmax) begin
`TRACE(1, ("%d: core%0d-scalar-tmem-softmax-start: grant=%b wid=%0d PC=0x%0h rd=%0d score=%0d prob=%0d tmask=%b (#%0d)\n",
$time, CORE_ID, scalar_tmem_grant, scalar_tmem_req_wid[0], scalar_tmem_req_pc[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]));
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,
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
`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,
scalar_tmem_wready, scalar_tmem_pc_r, scalar_tmem_rd_r));
end
if (scalar_tmem_commit_pending && scalar_tmem_pc_r == 32'h80000318) begin
`TRACE(1, ("%d: core%0d-scalar-tmem-softmax-commit-pending: grant_r=%b ready=%b valid0=%b wb0=%0d rd=%0d PC=0x%0h token=0x%0h\n",
$time, CORE_ID, scalar_tmem_grant_r, scalar_tmem_commit_ready,
scalar_tmem_commit_if[0].valid, scalar_tmem_commit_if[0].data.wb,
scalar_tmem_commit_if[0].data.rd, scalar_tmem_commit_if[0].data.PC,
scalar_tmem_softmax_token));
end
end
end
`endif
always @(*) begin
scalar_tmem_ren = scalar_tmem_grant_valid && scalar_tmem_grant_is_load;
scalar_tmem_wen = scalar_tmem_grant_valid && scalar_tmem_grant_is_store;
scalar_tmem_ren = (scalar_tmem_grant_valid && scalar_tmem_grant_is_load)
|| scalar_tmem_softmax_read_valid;
scalar_tmem_wen = (scalar_tmem_grant_valid && scalar_tmem_grant_is_store)
|| scalar_tmem_softmax_write_valid;
scalar_tmem_raddr = '0;
scalar_tmem_waddr = '0;
scalar_tmem_wdata = '0;
@@ -433,6 +559,14 @@ module VX_execute import VX_gpu_pkg::*; #(
& {(`NUM_THREADS * (`XLEN / 8)){scalar_tmem_grant_is_store}};
end
end
if (scalar_tmem_softmax_read_valid) begin
scalar_tmem_raddr = scalar_tmem_softmax_read_addr;
end
if (scalar_tmem_softmax_write_valid) begin
scalar_tmem_waddr = scalar_tmem_softmax_write_addr;
scalar_tmem_wdata = scalar_tmem_softmax_wdata;
scalar_tmem_mask = scalar_tmem_expand_tmask(scalar_tmem_tmask_r);
end
end
always @(posedge clk) begin
@@ -441,6 +575,13 @@ module VX_execute import VX_gpu_pkg::*; #(
scalar_tmem_load_pending <= 1'b0;
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_score_base_r <= '0;
scalar_tmem_softmax_p_base_r <= '0;
scalar_tmem_uuid_r <= '0;
scalar_tmem_wid_r <= '0;
scalar_tmem_tmask_r <= '0;
@@ -453,7 +594,7 @@ module VX_execute import VX_gpu_pkg::*; #(
scalar_tmem_grant_r <= scalar_tmem_grant;
scalar_tmem_load_pending <= scalar_tmem_grant_is_load;
scalar_tmem_store_pending <= scalar_tmem_grant_is_store;
scalar_tmem_commit_pending <= 1'b1;
scalar_tmem_commit_pending <= !scalar_tmem_grant_is_softmax;
scalar_tmem_rdata_valid_r <= 1'b0;
for (integer i = 0; i < `ISSUE_WIDTH; ++i) begin
if (scalar_tmem_grant[i]) begin
@@ -462,6 +603,19 @@ module VX_execute import VX_gpu_pkg::*; #(
scalar_tmem_tmask_r <= scalar_tmem_req_tmask[i];
scalar_tmem_pc_r <= scalar_tmem_req_pc[i];
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_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
end else if (scalar_tmem_load_pending && scalar_tmem_commit_pending && !scalar_tmem_rdata_valid_r) begin
@@ -469,6 +623,58 @@ module VX_execute import VX_gpu_pkg::*; #(
scalar_tmem_rdata_valid_r <= 1'b1;
end
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
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
end
if (scalar_tmem_commit_pending && (|(scalar_tmem_grant_r & scalar_tmem_commit_ready))) begin
scalar_tmem_grant_r <= '0;
scalar_tmem_load_pending <= 1'b0;