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 →
7.9 KiB228 linesraw
1#include "common.h"
2
3// OP: 0 - add, 1 - sub, 2 - mul, 3 - div
4constant short FC_bin_op [[function_constant(FC_BIN + 0)]];
5constant short FC_bin_f  [[function_constant(FC_BIN + 1)]];
6constant bool  FC_bin_rb [[function_constant(FC_BIN + 2)]];
7constant bool  FC_bin_cb [[function_constant(FC_BIN + 3)]];
8
9template <typename T0, typename T1, typename T>
10kernel void kernel_bin_fuse_impl(
11        constant ggml_metal_kargs_bin & args,
12        device const char * src0,
13        device const char * src1,
14        device       char * dst,
15        uint3   tgpig[[threadgroup_position_in_grid]],
16        ushort3 tpitg[[thread_position_in_threadgroup]],
17        ushort3   ntg[[threads_per_threadgroup]]) {
18#define FC_OP FC_bin_op
19#define FC_F  FC_bin_f
20#define FC_RB FC_bin_rb
21#define FC_CB FC_bin_cb
22
23    if (FC_RB) {
24        // row broadcast
25        const uint i0 = tgpig.y*args.ne00 + tgpig.x;
26        const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
27
28        device const T0 * src0_row = (device const T0 *) (src0);
29        device       T  * dst_row  = (device       T  *) (dst);
30
31        if (FC_F == 1) {
32            device const T1 * src1_row = (device const T1 *) (src1 + args.o1[0]);
33
34            if (FC_OP == 0) {
35                dst_row[i0] = src0_row[i0] + src1_row[i1];
36            }
37
38            if (FC_OP == 1) {
39                dst_row[i0] = src0_row[i0] - src1_row[i1];
40            }
41
42            if (FC_OP == 2) {
43                dst_row[i0] = src0_row[i0] * src1_row[i1];
44            }
45
46            if (FC_OP == 3) {
47                dst_row[i0] = src0_row[i0] / src1_row[i1];
48            }
49        } else {
50            T0 res = src0_row[i0];
51
52            if (FC_OP == 0) {
53                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
54                    res += ((device const T1 *) (src1 + args.o1[j]))[i1];
55                }
56            }
57
58            if (FC_OP == 1) {
59                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
60                    res -= ((device const T1 *) (src1 + args.o1[j]))[i1];
61                }
62            }
63
64            if (FC_OP == 2) {
65                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
66                    res *= ((device const T1 *) (src1 + args.o1[j]))[i1];
67                }
68            }
69
70            if (FC_OP == 3) {
71                FOR_UNROLL (short j = 0; j < FC_F; ++j) {
72                    res /= ((device const T1 *) (src1 + args.o1[j]))[i1];
73                }
74            }
75
76            dst_row[i0] = res;
77        }
78    } else {
79        const int i03 = tgpig.z;
80        const int i02 = tgpig.y;
81        const int i01 = tgpig.x;
82
83        if (i01 >= args.ne01) {
84            return;
85        }
86
87        const int i13 = i03%args.ne13;
88        const int i12 = i02%args.ne12;
89        const int i11 = i01%args.ne11;
90
91        device const T0 * src0_ptr = (device const T0 *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + args.offs);
92        device       T  * dst_ptr  = (device       T  *) (dst  + i03*args.nb3  + i02*args.nb2  + i01*args.nb1  + args.offs);
93
94        if (FC_F == 1) {
95            device const T1 * src1_ptr = (device const T1 *) (src1 + args.o1[0] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
96
97            for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
98                const int i10 = FC_CB ? i0%args.ne10 : i0;
99
100                if (FC_OP == 0) {
101                    dst_ptr[i0] = src0_ptr[i0] + src1_ptr[i10];
102                }
103
104                if (FC_OP == 1) {
105                    dst_ptr[i0] = src0_ptr[i0] - src1_ptr[i10];
106                }
107
108                if (FC_OP == 2) {
109                    dst_ptr[i0] = src0_ptr[i0] * src1_ptr[i10];
110                }
111
112                if (FC_OP == 3) {
113                    dst_ptr[i0] = src0_ptr[i0] / src1_ptr[i10];
114                }
115            }
116        } else {
117            device const T1 * src1_ptr[8];
118            FOR_UNROLL (short j = 0; j < FC_F; ++j) {
119                src1_ptr[j] = (device const T1 *) (src1 + args.o1[j] + i13*args.nb13 + i12*args.nb12 + i11*args.nb11);
120            }
121
122            for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
123                const int i10 = FC_CB ? i0%args.ne10 : i0;
124
125                T res = src0_ptr[i0];
126
127                if (FC_OP == 0) {
128                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
129                        res += src1_ptr[j][i10];
130                    }
131                }
132
133                if (FC_OP == 1) {
134                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
135                        res -= src1_ptr[j][i10];
136                    }
137                }
138
139                if (FC_OP == 2) {
140                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
141                        res *= src1_ptr[j][i10];
142                    }
143                }
144
145                if (FC_OP == 3) {
146                    FOR_UNROLL (short j = 0; j < FC_F; ++j) {
147                        res /= src1_ptr[j][i10];
148                    }
149                }
150
151                dst_ptr[i0] = res;
152            }
153        }
154    }
155
156#undef FC_OP
157#undef FC_F
158#undef FC_RB
159#undef FC_CB
160}
161
162typedef decltype(kernel_bin_fuse_impl<float, float, float>) kernel_bin_fuse_t;
163
164template [[host_name("kernel_bin_fuse_f32_f32_f32")]]   kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float,  float,  float>;
165template [[host_name("kernel_bin_fuse_f32_f32_f32_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<float4, float4, float4>;
166template [[host_name("kernel_bin_fuse_f16_f16_f16")]]   kernel kernel_bin_fuse_t kernel_bin_fuse_impl<half,   half,   half>;
167template [[host_name("kernel_bin_fuse_f16_f16_f16_4")]] kernel kernel_bin_fuse_t kernel_bin_fuse_impl<half4,  half4,  half4>;
168
169kernel void kernel_add_id(
170        constant ggml_metal_kargs_add_id & args,
171        device const char * src0,
172        device const char * src1,
173        device const char * src2,
174        device       char * dst,
175        uint3   tgpig[[threadgroup_position_in_grid]],
176        ushort3 tpitg[[thread_position_in_threadgroup]],
177        ushort3   ntg[[threads_per_threadgroup]]) {
178    const int i1 = tgpig.x;
179    const int i2 = tgpig.y;
180
181    const int i11 = *((device const int32_t *) (src2 + i1*sizeof(int32_t) + i2*args.nb21));
182
183    const size_t nb1 = args.ne0 * sizeof(float);
184    const size_t nb2 = args.ne1 * nb1;
185
186    device       float * dst_row  = (device       float *)((device char *)dst  +  i1*nb1       + i2*nb2);
187    device const float * src0_row = (device const float *)((device char *)src0 +  i1*args.nb01 + i2*args.nb02);
188    device const float * src1_row = (device const float *)((device char *)src1 + i11*args.nb11);
189
190    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
191        dst_row[i0] = src0_row[i0] + src1_row[i0];
192    }
193}
194
195template<typename T>
196kernel void kernel_repeat(
197        constant ggml_metal_kargs_repeat & args,
198        device const char * src0,
199        device       char * dst,
200        uint3   tgpig[[threadgroup_position_in_grid]],
201        ushort3 tpitg[[thread_position_in_threadgroup]],
202        ushort3   ntg[[threads_per_threadgroup]]) {
203    const int i3 = tgpig.z;
204    const int i2 = tgpig.y;
205    const int i1 = tgpig.x;
206
207    const int i03 = i3%args.ne03;
208    const int i02 = i2%args.ne02;
209    const int i01 = i1%args.ne01;
210
211    device const char * src0_ptr = src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01;
212    device       char * dst_ptr  = dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1;
213
214    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
215        const int i00 = i0%args.ne00;
216        *((device T *)(dst_ptr + i0*args.nb0)) = *((device T *)(src0_ptr + i00*args.nb00));
217    }
218}
219
220typedef decltype(kernel_repeat<float>) kernel_repeat_t;
221
222template [[host_name("kernel_repeat_f32")]] kernel kernel_repeat_t kernel_repeat<float>;
223template [[host_name("kernel_repeat_f16")]] kernel kernel_repeat_t kernel_repeat<half>;
224#if defined(GGML_METAL_HAS_BF16)
225template [[host_name("kernel_repeat_bf16")]] kernel kernel_repeat_t kernel_repeat<bfloat>;
226#endif
227template [[host_name("kernel_repeat_i32")]] kernel kernel_repeat_t kernel_repeat<int>;
228template [[host_name("kernel_repeat_i16")]] kernel kernel_repeat_t kernel_repeat<short>;