feat: add pipelined scalar tmem softmax unit

hw/dpi/float_dpi.cpp

@@ -12,6 +12,7 @@
 // limitations under the License.
 
 #include <stdio.h>
+#include <stdint.h>
 #include <math.h>
 #include <unordered_map>
 #include <vector>
@@ -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<uint32_t>(a_bits));
+  float b = bits_to_float(static_cast<uint32_t>(b_bits));
+  return static_cast<int>(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<uint32_t>(score_bits));
+  float row_max = bits_to_float(static_cast<uint32_t>(max_bits));
+  float accum = bits_to_float(static_cast<uint32_t>(accum_bits));
+  return static_cast<int>(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<uint32_t>(score_bits));
+  float row_max = bits_to_float(static_cast<uint32_t>(max_bits));
+  float denom = bits_to_float(static_cast<uint32_t>(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<int>(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<float>(frac) / 8.0f, -6);
+  }
+  return ldexpf(1.0f + static_cast<float>(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<uint8_t>(bits));
+    const float diff = fabsf(candidate - value);
+    if (diff < best_diff || (diff == best_diff && ((bits & 1u) == 0u))) {
+      best = static_cast<uint8_t>(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<uint32_t>(score_bits));
+  float row_max = bits_to_float(static_cast<uint32_t>(max_bits));
+  float denom = bits_to_float(static_cast<uint32_t>(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<int>(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;

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);

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]

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

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,13 @@ 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 [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  [`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 +376,10 @@ 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 [5:0] SCALAR_TMEM_SOFTMAX_ROW_LAST  = 6'd31;
+    localparam [5:0] SCALAR_TMEM_SOFTMAX_TILE_LAST = 6'd63;
 
     assign scalar_tmem_pending = |scalar_tmem_dispatch;
 
@@ -375,11 +394,56 @@ 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_reads_issued;
+    wire scalar_tmem_softmax_read_issue = scalar_tmem_softmax_read_valid && scalar_tmem_rready;
+    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
+                                          && 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_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;
@@ -390,15 +454,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 +485,41 @@ module VX_execute import VX_gpu_pkg::*; #(
         end
     endfunction
 
+`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_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_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_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
+                `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 +533,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 +549,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_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;
             scalar_tmem_wid_r <= '0;
             scalar_tmem_tmask_r <= '0;
@@ -453,7 +568,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 +577,15 @@ 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_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];
+                        end
                     end
                 end
             end else if (scalar_tmem_load_pending && scalar_tmem_commit_pending && !scalar_tmem_rdata_valid_r) begin
@@ -469,6 +593,21 @@ 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_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_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
                 scalar_tmem_grant_r <= '0;
                 scalar_tmem_load_pending <= 1'b0;

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