From dff1107bf5e8b4e38e2acb9c623b1286ac662ec4 Mon Sep 17 00:00:00 2001 From: Zhongdi LUO Date: Fri, 10 Jul 2026 08:23:22 +0000 Subject: [PATCH] feat: add scalar tmem softmax pipeline --- hw/dpi/float_dpi.cpp | 72 ++++++++++++ hw/dpi/float_dpi.vh | 4 + hw/rtl/VX_define.vh | 1 + hw/rtl/core/VX_decode.sv | 10 ++ hw/rtl/core/VX_execute.sv | 224 ++++++++++++++++++++++++++++++++++++-- 5 files changed, 302 insertions(+), 9 deletions(-) diff --git a/hw/dpi/float_dpi.cpp b/hw/dpi/float_dpi.cpp index 1f438ed9..197f48de 100644 --- a/hw/dpi/float_dpi.cpp +++ b/hw/dpi/float_dpi.cpp @@ -12,6 +12,7 @@ // limitations under the License. #include +#include #include #include #include @@ -38,6 +39,10 @@ extern "C" { void dpi_fdiv(bool enable, int dst_fmt, int64_t a, int64_t b, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags); void dpi_fsqrt(bool enable, int dst_fmt, int64_t a, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags); void dpi_fexp(bool enable, int dst_fmt, int64_t a, int64_t* result, svBitVecVal* fflags); + int dpi_f32_max(bool enable, int a_bits, int b_bits); + int dpi_softmax_exp_acc(bool enable, int score_bits, int max_bits, int accum_bits); + int dpi_softmax_prob_to_f16x2(bool enable, int score_bits, int max_bits, int denom_bits); + int dpi_softmax_prob_to_fp8e4m3x4(bool enable, int score_bits, int max_bits, int denom_bits); void dpi_ftoi(bool enable, int dst_fmt, int src_fmt, int64_t a, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags); void dpi_ftou(bool enable, int dst_fmt, int src_fmt, int64_t a, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags); @@ -100,6 +105,73 @@ inline uint32_t float_to_bits(float value) { return bits.u; } +int dpi_f32_max(bool enable, int a_bits, int b_bits) { + if (!enable) + return 0; + float a = bits_to_float(static_cast(a_bits)); + float b = bits_to_float(static_cast(b_bits)); + return static_cast(float_to_bits(a > b ? a : b)); +} + +int dpi_softmax_exp_acc(bool enable, int score_bits, int max_bits, int accum_bits) { + if (!enable) + return 0; + float score = bits_to_float(static_cast(score_bits)); + float row_max = bits_to_float(static_cast(max_bits)); + float accum = bits_to_float(static_cast(accum_bits)); + return static_cast(float_to_bits(accum + expf(score - row_max))); +} + +int dpi_softmax_prob_to_f16x2(bool enable, int score_bits, int max_bits, int denom_bits) { + if (!enable) + return 0; + float score = bits_to_float(static_cast(score_bits)); + float row_max = bits_to_float(static_cast(max_bits)); + float denom = bits_to_float(static_cast(denom_bits)); + float prob = denom == 0.0f ? 0.0f : expf(score - row_max) / denom; + half_float::half half_prob(prob); + uint16_t half_bits; + static_assert(sizeof(half_bits) == sizeof(half_prob), "unexpected half size"); + __builtin_memcpy(&half_bits, &half_prob, sizeof(half_bits)); + return static_cast(half_bits | (uint32_t(half_bits) << 16)); +} + +static float fp8e4m3_positive_to_float(uint8_t bits) { + const int exp = (bits >> 3) & 0xf; + const int frac = bits & 0x7; + if (exp == 0) { + return frac == 0 ? 0.0f : ldexpf(static_cast(frac) / 8.0f, -6); + } + return ldexpf(1.0f + static_cast(frac) / 8.0f, exp - 7); +} + +static uint8_t float_to_fp8e4m3_positive(float value) { + if (!(value > 0.0f)) + return 0; + uint8_t best = 0; + float best_diff = fabsf(value); + for (uint32_t bits = 1; bits < 0x80; ++bits) { + const float candidate = fp8e4m3_positive_to_float(static_cast(bits)); + const float diff = fabsf(candidate - value); + if (diff < best_diff || (diff == best_diff && ((bits & 1u) == 0u))) { + best = static_cast(bits); + best_diff = diff; + } + } + return best; +} + +int dpi_softmax_prob_to_fp8e4m3x4(bool enable, int score_bits, int max_bits, int denom_bits) { + if (!enable) + return 0; + float score = bits_to_float(static_cast(score_bits)); + float row_max = bits_to_float(static_cast(max_bits)); + float denom = bits_to_float(static_cast(denom_bits)); + float prob = denom == 0.0f ? 0.0f : expf(score - row_max) / denom; + uint32_t fp8 = float_to_fp8e4m3_positive(prob); + return static_cast(fp8 | (fp8 << 8) | (fp8 << 16) | (fp8 << 24)); +} + void dpi_fadd(bool enable, int dst_fmt, int64_t a, int64_t b, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags) { if (!enable) return; diff --git a/hw/dpi/float_dpi.vh b/hw/dpi/float_dpi.vh index 0644ac68..b214f067 100644 --- a/hw/dpi/float_dpi.vh +++ b/hw/dpi/float_dpi.vh @@ -27,6 +27,10 @@ import "DPI-C" function void dpi_fnmsub(input logic enable, input int dst_fmt, i import "DPI-C" function void dpi_fdiv(input logic enable, input int dst_fmt, input longint a, input longint b, input bit[2:0] frm, output longint result, output bit[4:0] fflags); import "DPI-C" function void dpi_fsqrt(input logic enable, input int dst_fmt, input longint a, input bit[2:0] frm, output longint result, output bit[4:0] fflags); import "DPI-C" function void dpi_fexp(input logic enable, input int dst_fmt, input longint a, output longint result, output bit[4:0] fflags); +import "DPI-C" function int dpi_f32_max(input logic enable, input int a_bits, input int b_bits); +import "DPI-C" function int dpi_softmax_exp_acc(input logic enable, input int score_bits, input int max_bits, input int accum_bits); +import "DPI-C" function int dpi_softmax_prob_to_f16x2(input logic enable, input int score_bits, input int max_bits, input int denom_bits); +import "DPI-C" function int dpi_softmax_prob_to_fp8e4m3x4(input logic enable, input int score_bits, input int max_bits, input int denom_bits); import "DPI-C" function void dpi_ftoi(input logic enable, input int dst_fmt, input int src_fmt, input longint a, input bit[2:0] frm, output longint result, output bit[4:0] fflags); import "DPI-C" function void dpi_ftou(input logic enable, input int dst_fmt, input int src_fmt, input longint a, input bit[2:0] frm, output longint result, output bit[4:0] fflags); diff --git a/hw/rtl/VX_define.vh b/hw/rtl/VX_define.vh index 8f03747c..60d1c695 100644 --- a/hw/rtl/VX_define.vh +++ b/hw/rtl/VX_define.vh @@ -210,6 +210,7 @@ `define INST_LSU_SD 4'b1011 // new for RV64I SD `define INST_LSU_TMEM_LD 4'b1100 `define INST_LSU_TMEM_ST 4'b1101 +`define INST_LSU_TMEM_SOFTMAX 4'b1110 `define INST_LSU_FENCE 4'b1111 `define INST_LSU_BITS 4 `define INST_LSU_FMT(op) op[2:0] diff --git a/hw/rtl/core/VX_decode.sv b/hw/rtl/core/VX_decode.sv index 1fd7f1c6..85a2c497 100644 --- a/hw/rtl/core/VX_decode.sv +++ b/hw/rtl/core/VX_decode.sv @@ -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 diff --git a/hw/rtl/core/VX_execute.sv b/hw/rtl/core/VX_execute.sv index 7a4e3d65..86e3d9ed 100644 --- a/hw/rtl/core/VX_execute.sv +++ b/hw/rtl/core/VX_execute.sv @@ -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;