char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Adrien Gallouëtggml : use a simple std::thread in AMX without OpenMP (#20074)66199c9f0

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.6 KiB115 linesraw
1#pragma once
2
3#include "ggml.h"
4#include "ggml-cpu-impl.h"
5
6#include <algorithm>
7#include <memory>
8#include <type_traits>
9
10#if defined(GGML_USE_OPENMP)
11#include <omp.h>
12#else
13#include <thread>
14#endif
15
16#define TILE_M 16
17#define TILE_N 16
18#define TILE_K 32
19#define VNNI_BLK 4
20
21#define AMX_BLK_SIZE 32
22
23#define TMM0 0
24#define TMM1 1
25#define TMM2 2
26#define TMM3 3
27#define TMM4 4
28#define TMM5 5
29#define TMM6 6
30#define TMM7 7
31
32// parallel routines
33template <typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0>
34inline T div_up(T x, T y) { return (x + y - 1) / y; }
35
36template <typename T>
37inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
38#if 0
39    // onednn partition pattern
40    T& n_my = n_end;
41    if (nth <= 1 || n == 0) {
42        n_start = 0;
43        n_my = n;
44    } else {
45        T n1 = div_up(n, nth);
46        T n2 = n1 - 1;
47        T T1 = n - n2 * nth;
48        n_my = ith < T1 ? n1 : n2;
49        n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2;
50    }
51    n_end += n_start;
52#else
53    // pytorch aten partition pattern
54    T n_my = div_up(n, nth);
55    n_start = ith * n_my;
56    n_end = std::min(n_start + n_my, n);
57#endif
58}
59
60template <typename func_t>
61inline void parallel_for(int n, const func_t & f) {
62    if (n <= 0) {
63        return;
64    }
65#if defined(GGML_USE_OPENMP)
66    #pragma omp parallel
67    {
68        int nth = omp_get_num_threads();
69        int ith = omp_get_thread_num();
70        int tbegin, tend;
71        balance211(n, nth, ith, tbegin, tend);
72        f(tbegin, tend);
73    }
74#else
75    int nth = std::thread::hardware_concurrency();
76    if (nth <= 1) {
77        f(0, n);
78        return;
79    }
80    if (nth > n) {
81        nth = n;
82    }
83    std::vector<std::thread> threads;
84    threads.reserve(nth);
85    for (int ith = 0; ith < nth; ++ith) {
86        threads.emplace_back([&f, n, ith, nth] {
87            int tbegin, tend;
88            balance211(n, nth, ith, tbegin, tend);
89            f(tbegin, tend);
90        });
91    }
92    for (auto & t : threads) {
93        t.join();
94    }
95#endif
96}
97
98template <typename func_t>
99inline void parallel_for_ggml(const ggml_compute_params * params, int n, const func_t & f) {
100    int tbegin, tend;
101    balance211(n, params->nth, params->ith, tbegin, tend);
102    f(tbegin, tend);
103}
104
105// quantized types that have AMX support
106inline bool qtype_has_amx_kernels(const enum ggml_type type) {
107    // TODO: fix padding for vnni format
108    return (type == GGML_TYPE_Q4_0) ||
109        (type == GGML_TYPE_Q4_1) ||
110        (type == GGML_TYPE_Q8_0) ||
111        (type == GGML_TYPE_Q4_K) ||
112        (type == GGML_TYPE_Q5_K) ||
113        (type == GGML_TYPE_Q6_K) ||
114        (type == GGML_TYPE_IQ4_XS);
115}