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.4 KiB75 linesraw
1#include "common.h"
2
3constant short FC_solve_tri_nsg [[function_constant(FC_SOLVE_TRI + 0)]];
4constant short FC_solve_tri_n   [[function_constant(FC_SOLVE_TRI + 1)]];
5constant short FC_solve_tri_k   [[function_constant(FC_SOLVE_TRI + 2)]];
6
7kernel void kernel_solve_tri_f32(
8        constant ggml_metal_kargs_solve_tri & args,
9        device   const char * src0,
10        device   const char * src1,
11        device         char * dst,
12        threadgroup    char * shmem [[threadgroup(0)]],
13        ushort3 tgpig[[threadgroup_position_in_grid]],
14        ushort  sgitg[[simdgroup_index_in_threadgroup]],
15        ushort  tiisg[[thread_index_in_simdgroup]],
16        ushort3   ntg[[threads_per_threadgroup]]) {
17    constexpr short NW = N_SIMDWIDTH;
18
19    const short NSG = FC_solve_tri_nsg;
20    const short N   = FC_solve_tri_n;
21    const short K   = FC_solve_tri_k;
22    const short NP  = PAD2(N, NW);
23
24    const int32_t i03 = tgpig.z;
25    const int32_t i02 = tgpig.y;
26    const int32_t i01 = tgpig.x*NSG + sgitg;
27
28    threadgroup float * sh0 = (threadgroup float *) shmem;
29
30    device const float * src0_ptr = (device const float *)(src0 + i02 * args.nb02 + i03 * args.nb03) + sgitg*N;
31    device const float * src1_ptr = (device const float *)(src1 + i02 * args.nb12 + i03 * args.nb13) + i01;
32    device       float * dst_ptr  = (device       float *)(dst  + i02 * args.nb2  + i03 * args.nb3)  + i01;
33
34    for (short rr = 0; rr < N; rr += NSG) {
35        threadgroup_barrier(mem_flags::mem_threadgroup);
36
37        {
38            threadgroup float * sh0_cur = sh0 + sgitg*NP;
39
40            for (short t = 0; t*NW < N; ++t) {
41                const short idx = t*NW + tiisg;
42                sh0_cur[idx] = src0_ptr[idx];
43            }
44
45            src0_ptr += NSG*N;
46        }
47
48        threadgroup_barrier(mem_flags::mem_threadgroup);
49
50        if (i01 >= args.ne10) {
51            continue;
52        }
53
54        for (short ir = 0; ir < NSG && rr + ir < N; ++ir) {
55            const short r = rr + ir;
56
57            threadgroup float * sh0_cur = sh0 + ir*NP;
58
59            float sum = 0.0f;
60
61            for (short t = 0; t*NW < r; ++t) {
62                const short idx = t*NW + tiisg;
63                sum += sh0_cur[idx] * dst_ptr[idx*K] * (idx < r);
64            }
65
66            sum = simd_sum(sum);
67
68            if (tiisg == 0) {
69                const float diag = sh0_cur[r];
70
71                dst_ptr[r*K] = (src1_ptr[r*K] - sum) / diag;
72            }
73        }
74    }
75}