math.hpp

#pragma once
#include <memory>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <random>
#include <ranges>
#include "gguf_types.hpp"
#if defined(__STDCPP_FLOAT16_T__)
#include <stdfloat>
#else
namespace std { using float16_t = _Float16; }
#endif

namespace misha::math {

struct math_expr_t {};

struct tensor_1d {
    template <std::derived_from<math_expr_t> Expr>
    void eval_expr(Expr expr);

    float* data = nullptr;
    uint32_t dims[1] = { 0 };
    uint32_t capacity = -1;

    auto resize(uint32_t new_size) -> tensor_1d& {
        if (capacity != uint32_t(-1) and capacity < new_size) [[unlikely]] {
            capacity = std::max({32u, new_size, capacity * 3 / 2});
            ::operator delete(data, std::align_val_t(32));
            data = (float*)::operator new(sizeof(float) * capacity, std::align_val_t(32));
        }
        dims[0] = new_size;
        return *this;
    }

    void free() noexcept {
        if (capacity != uint32_t(-1))
            ::operator delete(data, std::align_val_t(32));
    }

    auto operator[](uint32_t i) const noexcept -> float& {
        return data[i];
    }

    void operator=(std::derived_from<math_expr_t> auto expr) {
        eval_expr(expr);
    }
};

struct tensor_view {
    void* data = nullptr;
    uint32_t dims[2] = { 0, 1 };
    uint32_t ggml_type = gguf::GGML_TYPE_F32;
    uint32_t capacity = 0;

    auto operator[](uint32_t i) const noexcept -> tensor_1d {
        assert(ggml_type == gguf::GGML_TYPE_F32);
        return tensor_1d { .data = (float*)data + i * dims[0], .dims = { dims[0] } };
    }
};


struct trunk_q6_k {
    uint8_t ql[128];
    uint8_t qh[64];
    int8_t scales[16];
    std::float16_t d;
};

struct block_q8_0 {
    std::float16_t d;
    int8_t qs[32];
};

struct block_q4_0 {
    std::float16_t d;
    uint8_t qs[16];
};


// 将float32向量量化到q8_0
inline void quantize_q8_0(block_q8_0* __restrict dst, const float* __restrict src, size_t n) noexcept {
    for (size_t block_idx = 0; block_idx < n / 32; block_idx++) {
        const float* x = src + block_idx * 32;
        block_q8_0& block = dst[block_idx];

        float amax = 0.0f;
        #pragma omp simd reduction(max:amax)
        for (int i = 0; i < 32; i++)
            amax = std::max(amax, std::abs(x[i]));

        float d = amax / 127.0f;
        block.d = std::float16_t(d);
        d = d != 0.0f ? (1.0f / d) : 0.0f;

        #pragma omp simd
        for (int i = 0; i < 32; i++) {
            float val = std::round(x[i] * d);
            block.qs[i] = (int8_t)std::clamp(val, -127.0f, 127.0f);
        }
    }
}


// 将q6_K量化的一个trunk反量化回float32
inline void dequantize_q6_k_trunk(float* __restrict out, const trunk_q6_k* block) noexcept {
    const float master_scale = float(block->d);

    float subblock_scales[16];
    #pragma omp simd
    for (int i = 0; i < 16; i++)
        subblock_scales[i] = master_scale * block->scales[i];

    // 16 组 * 16 个元素划分为 4 次迭代，每次处理 4 组 * 16 个元素
    for (int batch_idx = 0; batch_idx < 4; batch_idx++) {
        const int half_offset = batch_idx / 2;
        const int quarter_offset = batch_idx % 2;
        float* __restrict batch_output = out + half_offset * 128 + quarter_offset * 16;

        const uint8_t* __restrict ql_0 = block->ql + half_offset * 64 + quarter_offset * 16;
        const uint8_t* __restrict ql_1 = block->ql + half_offset * 64 + 32 + quarter_offset * 16;
        const uint8_t* __restrict qh   = block->qh + half_offset * 32 + quarter_offset * 16;

        const float s1 = subblock_scales[half_offset * 8 + 0 + quarter_offset];
        const float s2 = subblock_scales[half_offset * 8 + 2 + quarter_offset];
        const float s3 = subblock_scales[half_offset * 8 + 4 + quarter_offset];
        const float s4 = subblock_scales[half_offset * 8 + 6 + quarter_offset];

        #pragma omp simd
        for (int i = 0; i < 16; i++) {
            int8_t q1 = int8_t((ql_0[i] & 0b0000'1111) | ((qh[i] & 0b0000'0011) << 4)) - 32;
            int8_t q2 = int8_t((ql_1[i] & 0b0000'1111) | ((qh[i] & 0b0000'1100) << 2)) - 32;
            int8_t q3 = int8_t((ql_0[i] >> 4) | (qh[i] & 0b0011'0000)) - 32;
            int8_t q4 = int8_t((ql_1[i] >> 4) | ((qh[i] & 0b1100'0000) >> 2)) - 32;

            batch_output[i] = s1 * q1;
            batch_output[i + 32] = s2 * q2;
            batch_output[i + 64] = s3 * q3;
            batch_output[i + 96] = s4 * q4;
        }
    }
}


// 从词向量矩阵中获取token对应的词向量
inline void get_embedding(float* __restrict x, const tensor_view& embeddings, uint32_t token, float scale = 1.0f) {
    assert(embeddings.ggml_type == gguf::GGML_TYPE_Q6_K);
    size_t trunk_count = embeddings.dims[0] / 256;
    const trunk_q6_k* embed_q6_k = (const trunk_q6_k*)embeddings.data + token * trunk_count;

    for (size_t i = 0; i < trunk_count; i++) {
        dequantize_q6_k_trunk(x + i * 256, embed_q6_k + i);
        #pragma omp simd aligned(x:32)
        for (size_t j = 0; j < 256; j++)
            x[i * 256 + j] *= scale;
    }
}


// 预计算RoPE旋转的频率表
inline auto compute_rope_freqs(size_t dim, float freq_base) -> std::unique_ptr<float[]> {
    auto freqs = std::make_unique<float[]>(dim / 2);
    for (size_t i = 0; i < dim / 2; i++)
        freqs[i] = 1.0f / std::pow(freq_base, 2.0f * i / dim);
    return freqs;
}


// 向向量中嵌入RoPE
inline void apply_rope(tensor_1d& x, size_t pos, const float* __restrict freqs) noexcept {
    for (size_t i = 0; i < x.dims[0] / 2; i++) {
        float angle = float(pos) * freqs[i];
        float cos_val = std::cos(angle);
        float sin_val = std::sin(angle);
        float x1 = x[i];
        float x2 = x[i + x.dims[0] / 2];
        x[i]  = x1 * cos_val - x2 * sin_val;
        x[i + x.dims[0] / 2] = x1 * sin_val + x2 * cos_val;
    }
}


// out = gelu(gate) * up
inline void geglu(float* out, size_t n, float* gate, const float* __restrict up) {
    auto gelu = [](float x) {
        constexpr float SQRT_2_OVER_PI = 0.7978845608f;
        constexpr float COEF = 0.044715f;
        float inner = SQRT_2_OVER_PI * (x + COEF * x * x * x);
        return 0.5f * x * (1.0f + std::tanh(inner));
    };

    #pragma omp simd aligned(out:32) aligned(gate:32) aligned(up:32)
    for (size_t i = 0; i < n; i++)
        out[i] = gelu(gate[i]) * up[i];
}


// output? = rmsnorm(vec, weights)
template <bool use_weight = true, bool use_output = false>
inline void rmsnorm(const tensor_1d& vec, float epsilon, const tensor_view& weights = {}, float* output = nullptr) noexcept {
    float ss = 0.0f;
    const float* __restrict w = use_weight ? (const float* __restrict)weights.data : nullptr;
    float* __restrict x = vec.data;

    #pragma omp simd reduction(+:ss) aligned(x:32)
    for (size_t i = 0; i < vec.dims[0]; i++)
        ss += x[i] * x[i];

    float scale = 1.0f / std::sqrt(ss / (float)vec.dims[0] + epsilon);

    #pragma omp simd aligned(x:32)
    for (size_t i = 0; i < vec.dims[0]; i++) {
        float weight = use_weight ? w[i] : 1.0f;
        if constexpr (use_output)
            output[i] = x[i] * scale * weight;
        else
            x[i] = x[i] * scale * weight;
    }
}

// x = softmax(x)
inline void softmax(const tensor_1d& x) {
    float max_val = std::numeric_limits<float>::lowest();
    #pragma omp simd reduction(max:max_val)
    for (uint32_t i = 0; i < x.dims[0]; i++)
        max_val = std::max(max_val, x[i]);

    float sum = 0.0f;
    #pragma omp simd reduction(+:sum)
    for (uint32_t i = 0; i < x.dims[0]; i++) {
        float exp_val = std::exp(x[i] - max_val);
        x[i] = exp_val;
        sum += exp_val;
    }

    const float inv_sum = 1.0f / sum;
    #pragma omp simd
    for (uint32_t i = 0; i < x.dims[0]; i++)
        x[i] *= inv_sum;
}


inline void vec_mat_mul(float* __restrict output, const tensor_1d& vec, const tensor_view& M) {
    assert(M.ggml_type == gguf::GGML_TYPE_F32);

    std::fill_n(output, M.dims[0], 0.0f);
    for (size_t i = 0; i < M.dims[1]; i++) {
        float weight = vec[i];
        auto m = (const float* __restrict)M.data + i * M.dims[0];
        #pragma omp simd
        for (size_t j = 0; j < M.dims[0]; j++)
            output[j] += weight * m[j];
    }
}


// float32向量点乘
inline auto vecdot(const float* __restrict a, const float* __restrict b, size_t n) noexcept -> float {
    [[assume(std::uintptr_t(a) % 32 == 0 and std::uintptr_t(b) % 32 == 0 and n % 32 == 0)]];

    float sum = 0.0f;
    #pragma omp simd reduction(+:sum) aligned(a:32) aligned(b:32)
    for (size_t i = 0; i < n; i++)
        sum += a[i] * b[i];
    return sum;
}

// float16向量与float32向量点乘
inline auto vecdot(const std::float16_t* __restrict a, const float* __restrict b, size_t n) noexcept -> float {
    [[assume(std::uintptr_t(a) % 32 == 0 and std::uintptr_t(b) % 32 == 0 and n % 32 == 0)]];

    float sum = 0.0f;
    #pragma omp simd
    for (size_t i = 0; i < n; i++)
        sum += a[i] * b[i];
    return sum;
}

// q4_0量化向量与q8_0量化向量点乘
inline auto vecdot(const block_q4_0* __restrict a, const block_q8_0* __restrict b, size_t n) noexcept -> float {
    float sum = 0.0f;
    for (size_t block_idx = 0; block_idx < n / 32; block_idx++) {
        const block_q4_0& block_a = a[block_idx];
        const block_q8_0& block_b = b[block_idx];

        int int_sum = 0;
        #pragma omp simd
        for (int i = 0; i < 16; i++) {
            uint8_t a_byte = block_a.qs[i];
            int8_t q_a_0 = (a_byte & 0b0000'1111) - 8;
            int8_t q_a_1 = (a_byte >> 4) - 8;
            int8_t q_b_0 = block_b.qs[i];
            int8_t q_b_1 = block_b.qs[i + 16];
            int_sum += q_a_0 * q_b_0 + q_a_1 * q_b_1;
        }

        sum += (float)int_sum * block_a.d * block_b.d;
    }
    return sum;
}


// q6_k量化向量与q8_0量化向量点乘
inline auto vecdot(const trunk_q6_k* __restrict a, const block_q8_0* __restrict b, size_t n) noexcept -> float {
    float total_sum = 0.0f;
    const size_t num_blocks = n / 256;

    for (size_t block_idx = 0; block_idx < num_blocks; block_idx++) {
        const trunk_q6_k& block_a = a[block_idx];
        const block_q8_0* block_b = b + block_idx * 8;
        const float d_a = float(block_a.d);

        float block_sum = 0.0f;
        for (int half_offset = 0; half_offset < 2; half_offset++) {
            const uint8_t* __restrict ql_base = block_a.ql + half_offset * 64;
            const uint8_t* __restrict qh_base = block_a.qh + half_offset * 32;
            const block_q8_0* __restrict b_curr = block_b + half_offset * 4;
            const int8_t* __restrict scales_curr = block_a.scales + half_offset * 8;

            int term[4] = {0, 0, 0, 0};
            for (int quarter_offset = 0; quarter_offset < 2; quarter_offset++) {
                const uint8_t* __restrict ql_0 = ql_base + quarter_offset * 16;
                const uint8_t* __restrict ql_1 = ql_base + 32 + quarter_offset * 16;
                const uint8_t* __restrict qh   = qh_base + quarter_offset * 16;

                const int8_t s1 = scales_curr[0 + quarter_offset];
                const int8_t s2 = scales_curr[2 + quarter_offset];
                const int8_t s3 = scales_curr[4 + quarter_offset];
                const int8_t s4 = scales_curr[6 + quarter_offset];

                const int8_t* __restrict b1 = b_curr[0].qs + quarter_offset * 16;
                const int8_t* __restrict b2 = b_curr[1].qs + quarter_offset * 16;
                const int8_t* __restrict b3 = b_curr[2].qs + quarter_offset * 16;
                const int8_t* __restrict b4 = b_curr[3].qs + quarter_offset * 16;

                int sum1 = 0, sum2 = 0, sum3 = 0, sum4 = 0;

                #pragma omp simd reduction(+:sum1, sum2, sum3, sum4)
                for (int i = 0; i < 16; i++) {
                    int q1 = int((ql_0[i] & 0x0F) | ((qh[i] & 0x03) << 4)) - 32;
                    int q2 = int((ql_1[i] & 0x0F) | ((qh[i] & 0x0C) << 2)) - 32;
                    int q3 = int((ql_0[i] >> 4)   | (qh[i] & 0x30))        - 32;
                    int q4 = int((ql_1[i] >> 4)   | ((qh[i] & 0xC0) >> 2)) - 32;

                    sum1 += q1 * b1[i];
                    sum2 += q2 * b2[i];
                    sum3 += q3 * b3[i];
                    sum4 += q4 * b4[i];
                }

                term[0] += sum1 * s1;
                term[1] += sum2 * s2;
                term[2] += sum3 * s3;
                term[3] += sum4 * s4;
            }

            block_sum += (float)term[0] * (d_a * float(b_curr[0].d)) +
                         (float)term[1] * (d_a * float(b_curr[1].d)) +
                         (float)term[2] * (d_a * float(b_curr[2].d)) +
                         (float)term[3] * (d_a * float(b_curr[3].d));
        }

        total_sum += block_sum;
    }
    return total_sum;
}

// 矩阵乘法：Y.T = X.T * M.T
inline auto matmul(float* __restrict Y, const tensor_view& X, const tensor_view& M) noexcept {
    thread_local auto x_q8_0 = std::unique_ptr<block_q8_0[]>();
    thread_local size_t x_q8_0_size = 0;

    auto matmul_float = [&]<class T> {
        for (uint32_t i = 0; i < X.dims[1]; i++) {
            float* __restrict y = Y + i * M.dims[1];
            float* __restrict x = (float*)X.data + i * X.dims[0];
            for (uint32_t j = 0; j < M.dims[1]; j++) {
                T* __restrict m = (T*)M.data + j * M.dims[0];
                y[j] = vecdot(m, x, X.dims[0]);
            }
        }
    };

    auto matmul_quantize = [&]<class T, size_t block_size> {
        if (X.dims[0] > x_q8_0_size) [[unlikely]] {
            x_q8_0_size = X.dims[0];
            x_q8_0.reset(new block_q8_0[x_q8_0_size / 32]);
        }

        for (uint32_t i = 0; i < X.dims[1]; i++) {
            float* __restrict y = Y + i * M.dims[1];
            float* __restrict x = (float*)X.data + i * X.dims[0];
            quantize_q8_0(x_q8_0.get(), x, X.dims[0]);

            for (uint32_t j = 0; j < M.dims[1]; j++) {
                T* __restrict m = (T*)M.data + j * (M.dims[0] / block_size);
                y[j] = vecdot(m, x_q8_0.get(), X.dims[0]);
            }
        }
    };

    switch(M.ggml_type) {
        case gguf::GGML_TYPE_F32: matmul_float.operator()<float>(); break;
        case gguf::GGML_TYPE_F16: matmul_float.operator()<std::float16_t>(); break;
        case gguf::GGML_TYPE_Q4_0: matmul_quantize.operator()<block_q4_0, 32>(); break;
        case gguf::GGML_TYPE_Q6_K: matmul_quantize.operator()<trunk_q6_k, 256>(); break;
        default: assert(false && "unsupported GGML type in matmul"); break;
    }
}


// 矩阵乘法：y = x * M
inline auto matmul(float* y, float* x, const tensor_view& M) noexcept {
    matmul(y, tensor_view { .data = x, .dims = { M.dims[0] , 1 } }, M);
}


// 对一组逻辑值进行 temperature + Top-P 采样，返回采样的逻辑值对应索引
auto inline sample(std::span<const float> logits, std::mt19937& random_generator, float temperature, float top_p) -> size_t {
    if (temperature <= 0.01f)
        return std::distance(logits.begin(), std::ranges::max_element(logits));

    const float max_logit = std::ranges::max(logits);

    thread_local auto indexed_weights = std::vector<std::pair<float, int>>();
    indexed_weights.clear();
    float total_weight = 0.0f;

    for (size_t i = 0; i < logits.size(); i++) {
        float weight = std::exp((logits[i] - max_logit) / temperature);
        indexed_weights.emplace_back(weight, i);
        total_weight += weight;
    }

    std::ranges::make_heap(indexed_weights);

    // 不断将大顶堆中权重最大的元素移到末尾，直到累加权重 weight_cumulative 达到 top_p * total_weight
    float weight_cumulative = 0.0f;
    for (size_t i = 0; i < indexed_weights.size(); i++) {
        const auto [top_weight, _] = indexed_weights[0];
        std::pop_heap(indexed_weights.begin(), indexed_weights.end() - i);
        if ((weight_cumulative += top_weight) >= top_p * total_weight)
            break;
    }

    // 在区间 [0, weight_cumulative) 随机取一个值 target_weight
    auto distribution = std::uniform_real_distribution<float>(0.0f, weight_cumulative);
    const float target_weight = distribution(random_generator);

    // 将权重从大到小累加，直到累加权重 weight_cumulative 达到 target_weight
    weight_cumulative = 0.0f;
    for (auto [weight, index] : indexed_weights | std::views::reverse) {
        weight_cumulative += weight;
        if (weight_cumulative >= target_weight)
            return index;
    }

    std::unreachable();
}

}