char-fork/llama.cpp

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

YiChen Lvmetal: per-op source split + parallel compile (#26561)b615f5b4b

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
2.8 KiB69 linesraw
1#include "common.h"
2
3template<uint32_t ttype>
4bool _ggml_vec_tri_cmp(const int i, const int r);
5
6template<>
7bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER */ 3>(const int i, const int r) {
8    return i < r;
9}
10
11template<>
12bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_LOWER_DIAG */ 2>(const int i, const int r) {
13    return i <= r;
14}
15
16template<>
17bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER */ 1>(const int i, const int r) {
18    return i > r;
19}
20
21template<>
22bool _ggml_vec_tri_cmp</* GGML_TRI_TYPE_UPPER_DIAG */ 0>(const int i, const int r) {
23    return i >= r;
24}
25
26template<typename T, int ttype>
27kernel void kernel_tri(
28        constant ggml_metal_kargs_tri & args,
29        device const char * src0,
30        device const char * dst,
31        uint3   tgpig[[threadgroup_position_in_grid]],
32        ushort3 tpitg[[thread_position_in_threadgroup]],
33        ushort3   ntg[[threads_per_threadgroup]]) {
34    const int i3 = tgpig.z;
35    const int i2 = tgpig.y;
36    const int i1 = tgpig.x;
37
38    if (i3 >= args.ne03 || i2 >= args.ne02 || i1 >= args.ne01) {
39        return;
40    }
41
42    device const T * src_row = (device const T *) ((device const char *) src0 + i1*args.nb01 + i2*args.nb02 + i3*args.nb03);
43    device       T * dst_row = (device       T *) ((device       char *) dst  + i1*args.nb1  + i2*args.nb2  + i3*args.nb3);
44
45    // Each thread is a single element of the row if ne00 < max threads per
46    // threadgroup, so this will loop once for each index that this thread is
47    // responsible for
48    for (int64_t i0 = tpitg.x; i0 < args.ne00; i0 += ntg.x) {
49        // Use the comparison as a mask for branchless
50        dst_row[i0] = static_cast<T>(_ggml_vec_tri_cmp<ttype>(i0, i1)) * src_row[i0];
51    }
52}
53
54typedef decltype(kernel_tri<float, 0>) kernel_tri_t;
55
56template [[host_name("kernel_tri_f32_0")]] kernel kernel_tri_t kernel_tri<float, 0>;
57template [[host_name("kernel_tri_f32_1")]] kernel kernel_tri_t kernel_tri<float, 1>;
58template [[host_name("kernel_tri_f32_2")]] kernel kernel_tri_t kernel_tri<float, 2>;
59template [[host_name("kernel_tri_f32_3")]] kernel kernel_tri_t kernel_tri<float, 3>;
60template [[host_name("kernel_tri_f16_0")]] kernel kernel_tri_t kernel_tri<half, 0>;
61template [[host_name("kernel_tri_f16_1")]] kernel kernel_tri_t kernel_tri<half, 1>;
62template [[host_name("kernel_tri_f16_2")]] kernel kernel_tri_t kernel_tri<half, 2>;
63template [[host_name("kernel_tri_f16_3")]] kernel kernel_tri_t kernel_tri<half, 3>;
64#if defined(GGML_METAL_HAS_BF16)
65template [[host_name("kernel_tri_bf16_0")]] kernel kernel_tri_t kernel_tri<bfloat, 0>;
66template [[host_name("kernel_tri_bf16_1")]] kernel kernel_tri_t kernel_tri<bfloat, 1>;
67template [[host_name("kernel_tri_bf16_2")]] kernel kernel_tri_t kernel_tri<bfloat, 2>;
68template [[host_name("kernel_tri_bf16_3")]] kernel kernel_tri_t kernel_tri<bfloat, 3>;
69#endif