feat: add scalar tmem softmax pipeline
This commit is contained in:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user