feat: add scalar tmem softmax pipeline

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

View File

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