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
6.8 KiB179 linesraw
1#include "common.h"
2
3constant bool FC_upscale_aa [[function_constant(FC_UPSCALE + 0)]];
4
5kernel void kernel_upscale_nearest_f32(
6    constant ggml_metal_kargs_upscale & args,
7    device  const char * src0,
8    device        char * dst,
9    uint3 tgpig[[threadgroup_position_in_grid]],
10    uint3 tpitg[[thread_position_in_threadgroup]],
11    uint3   ntg[[threads_per_threadgroup]]) {
12
13    const int64_t i3 = tgpig.z;
14    const int64_t i2 = tgpig.y;
15    const int64_t i1 = tgpig.x;
16
17    const int64_t i03 = i3/args.sf3;
18    const int64_t i02 = i2/args.sf2;
19    const int64_t i01 = i1/args.sf1;
20
21    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
22        const int64_t i00 = i0/args.sf0;
23
24        device const float * src0_ptr = (device const float *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
25        device       float * dst_ptr  = (device       float *) (dst  +  i3*args.nb3  +  i2*args.nb2  +  i1*args.nb1  +  i0*args.nb0);
26
27        dst_ptr[0] = src0_ptr[0];
28    }
29}
30
31static inline float bilinear_tri(float x) {
32    return MAX(0.0f, 1.0f - fabs(x));
33}
34
35kernel void kernel_upscale_bilinear_f32(
36    constant ggml_metal_kargs_upscale & args,
37    device  const char * src0,
38    device        char * dst,
39    uint3 tgpig[[threadgroup_position_in_grid]],
40    uint3 tpitg[[thread_position_in_threadgroup]],
41    uint3   ntg[[threads_per_threadgroup]]) {
42
43    const int64_t i3 = tgpig.z;
44    const int64_t i2 = tgpig.y;
45    const int64_t i1 = tgpig.x;
46
47    const int64_t i03 = i3 / args.sf3;
48    const int64_t i02 = i2 / args.sf2;
49
50    const float   f01  = ((float)i1 + args.poffs) / args.sf1 - args.poffs;
51    const int64_t i01  = MAX(0, MIN(args.ne01 - 1, (int64_t)floor(f01)));
52    const int64_t i01p = MAX(0, MIN(args.ne01 - 1, i01 + 1));
53    const float   fd1  = MAX(0.0f, MIN(1.0f, f01 - (float)i01));
54
55    src0 += i03*args.nb03 + i02*args.nb02;
56
57    device float * dst_ptr = (device float *)(dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1);
58
59    if (FC_upscale_aa) {
60        const float support0  = MAX(1.0f, 1.0f / args.sf0);
61        const float invscale0 = 1.0f / support0;
62        const float support1  = MAX(1.0f, 1.0f / args.sf1);
63        const float invscale1 = 1.0f / support1;
64
65        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
66            const float f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
67
68            int64_t x_min = MAX((int64_t)0, (int64_t)floor(f00 - support0 + args.poffs));
69            int64_t x_max = MIN(args.ne00,  (int64_t)ceil (f00 + support0 + args.poffs));
70
71            int64_t y_min = MAX((int64_t)0, (int64_t)floor(f01 - support1 + args.poffs));
72            int64_t y_max = MIN(args.ne01,  (int64_t)ceil (f01 + support1 + args.poffs));
73
74            float sum = 0.0f;
75            float wsum = 0.0f;
76
77            for (int64_t sy = y_min; sy < y_max; ++sy) {
78                const float wy = MAX(0.0f, 1.0f - fabs((float)sy - f01) * invscale1);
79                for (int64_t sx = x_min; sx < x_max; ++sx) {
80                    const float wx = MAX(0.0f, 1.0f - fabs((float)sx - f00) * invscale0);
81                    const float w  = wx * wy;
82                    device const float * src_ptr = (device const float *)(src0 + sy*args.nb01 + sx*args.nb00);
83                    sum  += (*src_ptr) * w;
84                    wsum += w;
85                }
86            }
87
88            const float v = (wsum > 0.0f) ? (sum / wsum) : 0.0f;
89            dst_ptr[i0] = v;
90        }
91    } else {
92        for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
93            const float   f00  = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
94            const int64_t i00  = MAX(0, MIN(args.ne00 - 1, (int64_t)floor(f00)));
95            const int64_t i00p = MAX(0, MIN(args.ne00 - 1, i00 + 1));
96            const float   fd0  = MAX(0.0f, MIN(1.0f, f00 - (float)i00));
97
98            device const float * src00 = (device const float *)(src0 + i01*args.nb01  + i00*args.nb00);
99            device const float * src10 = (device const float *)(src0 + i01*args.nb01  + i00p*args.nb00);
100            device const float * src01 = (device const float *)(src0 + i01p*args.nb01 + i00*args.nb00);
101            device const float * src11 = (device const float *)(src0 + i01p*args.nb01 + i00p*args.nb00);
102
103            const float v =
104                (*src00) * (1.0f - fd0) * (1.0f - fd1) +
105                (*src10) * fd0          * (1.0f - fd1) +
106                (*src01) * (1.0f - fd0) * fd1 +
107                (*src11) * fd0          * fd1;
108
109            dst_ptr[i0] = v;
110        }
111    }
112}
113
114static inline float bicubic_weight1(float x) {
115    const float a = -0.75f;
116    return ((a + 2) * x - (a + 3)) * x * x + 1;
117}
118
119static inline float bicubic_weight2(float x) {
120    const float a = -0.75f;
121    return ((a * x - 5 * a) * x + 8 * a) * x - 4 * a;
122}
123
124kernel void kernel_upscale_bicubic_f32(
125    constant ggml_metal_kargs_upscale & args,
126    device  const char * src0,
127    device        char * dst,
128    uint3 tgpig[[threadgroup_position_in_grid]],
129    uint3 tpitg[[thread_position_in_threadgroup]],
130    uint3   ntg[[threads_per_threadgroup]]) {
131
132    const int64_t i3 = tgpig.z;
133    const int64_t i2 = tgpig.y;
134    const int64_t i1 = tgpig.x;
135
136    const int64_t i03 = i3 / args.sf3;
137    const int64_t i02 = i2 / args.sf2;
138
139    const float   f01 = ((float)i1 + args.poffs) / args.sf1 - args.poffs;
140    const int64_t i01 = (int64_t)floor(f01);
141    const float   fd1 = f01 - (float)i01;
142
143    const float w_y0 = bicubic_weight2(fd1 + 1.0f);
144    const float w_y1 = bicubic_weight1(fd1);
145    const float w_y2 = bicubic_weight1(1.0f - fd1);
146    const float w_y3 = bicubic_weight2(2.0f - fd1);
147
148    const device char * src_slice = src0 + i03 * args.nb03 + i02 * args.nb02;
149
150    device float * dst_ptr = (device float *)(dst + i3 * args.nb3 + i2 * args.nb2 + i1 * args.nb1);
151
152    for (int i0 = tpitg.x; i0 < args.ne0; i0 += ntg.x) {
153        const float   f00 = ((float)i0 + args.poffs) / args.sf0 - args.poffs;
154        const int64_t i00 = (int64_t)floor(f00);
155        const float   fd0 = f00 - (float)i00;
156
157        const float w_x0 = bicubic_weight2(fd0 + 1.0f);
158        const float w_x1 = bicubic_weight1(fd0);
159        const float w_x2 = bicubic_weight1(1.0f - fd0);
160        const float w_x3 = bicubic_weight2(2.0f - fd0);
161
162        float sum = 0.0f;
163
164        for (int dy = -1; dy <= 2; ++dy) {
165            const int64_t iy = MAX(0, MIN(args.ne01 - 1, i01 + dy));
166            const float wy = (dy == -1) ? w_y0 : (dy == 0) ? w_y1 : (dy == 1) ? w_y2 : w_y3;
167
168            for (int dx = -1; dx <= 2; ++dx) {
169                const int64_t ix = MAX(0, MIN(args.ne00 - 1, i00 + dx));
170                const float wx = (dx == -1) ? w_x0 : (dx == 0) ? w_x1 : (dx == 1) ? w_x2 : w_x3;
171
172                device const float * src_ptr = (device const float *)(src_slice + iy * args.nb01 + ix * args.nb00);
173                sum += (*src_ptr) * wx * wy;
174            }
175        }
176
177        dst_ptr[i0] = sum;
178    }
179}