mini-llm-inference-engine/gemm_layout_test.cpp at main · JohnScheuer/mini-llm-inference-engine · 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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
#include <vector>
#include <iostream>
#include <cmath>
#include <algorithm>
#include <random>

//////////////////////////////////////////////////////////////
// Row-major GEMM (M x K) × (K)
//////////////////////////////////////////////////////////////
void gemm_row_major(
    const std::vector<float>& W,
    const std::vector<float>& x,
    std::vector<float>& y,
    int M,
    int K)
{
    for(int i=0;i<M;i++)
    {
        float acc = 0.f;

        for(int j=0;j<K;j++)
            acc += W[i*K + j] * x[j];

        y[i] = acc;
    }
}

//////////////////////////////////////////////////////////////
// Convert Row → Col major
//////////////////////////////////////////////////////////////
std::vector<float> to_col_major(
    const std::vector<float>& W,
    int M,
    int K)
{
    std::vector<float> W_col(M*K);

    for(int i=0;i<M;i++)
        for(int j=0;j<K;j++)
            W_col[j*M + i] = W[i*K + j];

    return W_col;
}

//////////////////////////////////////////////////////////////
// Col-major GEMM (M x K) × (K)
//////////////////////////////////////////////////////////////
void gemm_col_major(
    const std::vector<float>& W_col,
    const std::vector<float>& x,
    std::vector<float>& y,
    int M,
    int K)
{
    for(int i=0;i<M;i++)
    {
        float acc = 0.f;

        for(int j=0;j<K;j++)
            acc += W_col[j*M + i] * x[j];

        y[i] = acc;
    }
}

//////////////////////////////////////////////////////////////
// MAIN
//////////////////////////////////////////////////////////////
int main()
{
    const int M = 128;
    const int K = 128;

    std::cout << "--- GEMM Layout Test ---\n";

    std::vector<float> W(M*K);
    std::vector<float> x(K);
    std::vector<float> y_row(M);
    std::vector<float> y_col(M);

    std::mt19937 rng(42);
    std::uniform_real_distribution<float> dist(-1.f, 1.f);

    for(int i=0;i<M*K;i++)
        W[i] = dist(rng);

    for(int i=0;i<K;i++)
        x[i] = dist(rng);

    gemm_row_major(W, x, y_row, M, K);

    auto W_col = to_col_major(W, M, K);

    gemm_col_major(W_col, x, y_col, M, K);

    float max_diff = 0.f;
    float mean_diff = 0.f;

    for(int i=0;i<M;i++)
    {
        float diff = std::abs(y_row[i] - y_col[i]);
        max_diff = std::max(max_diff, diff);
        mean_diff += diff;
    }

    mean_diff /= M;

    std::cout << "Max abs diff:  " << max_diff << "\n";
    std::cout << "Mean abs diff: " << mean_diff << "\n";

    return 0;
}