gemma4.cpp/math_expr.hpp at main · SyrieYume/gemma4.cpp · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
// 通过表达式模板实现类似 Y = X * M 的写法，将表达式映射到对应计算函数
#pragma once
#include "math.hpp"

namespace misha::math {
    inline auto get_embedding(const tensor_view& embeddings, uint32_t token_id, float scale) noexcept {
        struct expr_t : math_expr_t { const tensor_view& embeddings; uint32_t token_id; float scale; };
        return expr_t { {}, embeddings, token_id, scale };
    }

    inline auto rmsnorm_to(tensor_1d& x, float epsilon, const tensor_view& weights = {}) noexcept {
        struct expr_t : math_expr_t { tensor_1d& x; float epsilon; const tensor_view& weights; };
        return expr_t { {}, x, epsilon, weights };
    }

    inline auto gelu(const tensor_1d& x) noexcept {
        struct expr_t : math_expr_t { const tensor_1d& x; };
        return expr_t { {}, x };
    }

    inline auto vec_mat_mul(const tensor_1d& vec, const tensor_view& M) noexcept {
        struct expr_t : math_expr_t { const tensor_1d& vec; const tensor_view& M; };
        return expr_t { {}, vec, M };
    }

    inline auto operator+(const tensor_1d& left_, const tensor_1d& right_) noexcept {
        struct expr_t : math_expr_t { const tensor_1d& left, right; };
        return expr_t { {}, left_, right_ };
    }

    inline auto operator*(const tensor_1d& val_, const tensor_view& M) noexcept {
        struct expr_t : math_expr_t { decltype(val_)& val; const tensor_view& M; };
        return expr_t { {}, val_, M };
    }

    inline auto operator*(decltype(gelu(tensor_1d())) left_, const tensor_1d& right_) noexcept {
        struct expr_t : math_expr_t { decltype(left_) left; decltype(right_) right; };
        return expr_t { {}, left_, right_ };
    }

    inline auto operator*(decltype(tensor_1d() + tensor_1d()) left_, float right_) noexcept {
        struct expr_t : math_expr_t { decltype(left_) left; float right; };
        return expr_t { {}, left_, right_ };
    }

    inline auto operator*(decltype(tensor_1d() * tensor_view()) left_, float right_) noexcept {
        struct expr_t : math_expr_t { decltype(left_) left; float right; };
        return expr_t { {}, left_, right_ };
    }


    template <> void tensor_1d::eval_expr(decltype(get_embedding({}, 0, 0.0f)) expr) {
        resize(expr.embeddings.dims[0]);
        get_embedding(data, expr.embeddings, expr.token_id, expr.scale);
    }

    template <> void tensor_1d::eval_expr(decltype(tensor_1d() * tensor_view()) expr) {
        resize(expr.M.dims[1]);
        matmul(data, expr.val.data, expr.M);
    }

    template <> void tensor_1d::eval_expr(decltype((tensor_1d() + tensor_1d()) * 1.0f) expr) {
        resize(expr.left.left.dims[0]);
        float* out = data;
        const float* a = expr.left.left.data;
        const float* __restrict b = expr.left.right.data;
        float scale = expr.right;

        #pragma omp simd aligned(out:32) aligned(a:32) aligned(b:32)
        for (size_t i = 0; i < dims[0]; i++)
            out[i] = (a[i] + b[i]) * scale;
    }

    template <> void tensor_1d::eval_expr(decltype(rmsnorm_to(std::declval<tensor_1d&>(), 0.0f)) expr) {
        resize(expr.x.dims[0]);
        rmsnorm<true, true>(expr.x, expr.epsilon, expr.weights, data);
    }

    template <> void tensor_1d::eval_expr(decltype((tensor_1d() * tensor_view()) * 1.0f) expr) {
        resize(expr.left.M.dims[1]);
        matmul(data, expr.left.val.data, expr.left.M);
        float* out = data;

        #pragma omp simd aligned(out:32)
        for (size_t i = 0; i < dims[0]; i++)
            out[i] *= expr.right;
    }

    template <> void tensor_1d::eval_expr(decltype(gelu(tensor_1d()) * tensor_1d()) expr) {
        resize(expr.left.x.dims[0]);
        geglu(data, dims[0], expr.left.x.data, expr.right.data);
    }

    template <> void tensor_1d::eval_expr(decltype(vec_mat_mul(tensor_1d(), tensor_view())) expr) {
        resize(expr.M.dims[0]);
        vec_mat_mul(data, expr.vec, expr.M);
    }
}