char-fork/llama.cpp

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

Nikhil Jain[ggml-webgpu] Implement 2D workgroups for scale, binary, and unary ops (#24044)3b3da01dc

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
6.1 KiB213 linesraw
1#ifdef TYPE_F16
2enable f16;
3#define TYPE f16
4#else
5#define TYPE f32
6#endif
7
8@group(0) @binding(0)
9var<storage, read_write> src: array<TYPE>;
10
11#ifndef INPLACE
12@group(0) @binding(1)
13var<storage, read_write> dst: array<TYPE>;
14#define PARAMS_BINDING 2
15#else
16#define PARAMS_BINDING 1
17#endif
18
19struct Params {
20    ne: u32,            // total number of elements
21    offset_src: u32,    // in elements
22    offset_dst: u32,    // in elements
23
24    // Strides (in elements)
25    stride_src0: u32,
26    stride_src1: u32,
27    stride_src2: u32,
28    stride_src3: u32,
29
30    // Logical shapes
31    ne0: u32,
32    ne1: u32,
33    ne2: u32,
34#ifdef CLAMP
35    clamp_min: f32,
36    clamp_max: f32,
37#endif
38#ifdef FILL
39    fill_val: f32,
40#endif
41#ifdef XIELU
42    alpha_n: f32,
43    alpha_p: f32,
44    beta: f32,
45    eps: f32,
46#endif
47
48};
49
50@group(0) @binding(PARAMS_BINDING)
51var<uniform> params: Params;
52
53fn erf_approx(x: TYPE) -> TYPE {
54    let x_f32 = f32(x);
55    let s = select(-1.0, 1.0, x_f32 >= 0.0);
56    let ax = abs(x_f32);
57
58    let t = 1.0 / (1.0 + 0.3275911 * ax);
59
60    let y = 1.0 -
61        (((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t
62            - 0.284496736) * t + 0.254829592) * t) *
63        exp(-ax * ax);
64
65    return TYPE(s * y);
66}
67
68@compute @workgroup_size(WG_SIZE)
69fn main(@builtin(global_invocation_id) gid: vec3<u32>,
70    @builtin(num_workgroups)       num_wg:  vec3<u32>) {
71    let threads_per_group = u32(WG_SIZE);
72    let flat_i = gid.x + (num_wg.x * threads_per_group) * gid.y;
73    if (flat_i >= params.ne) {
74        return;
75    }
76    var i = flat_i;
77    let ne2 = params.ne2;
78#ifdef DIAG
79    let ne1 = params.ne0;
80#else
81    let ne1 = params.ne1;
82#endif
83    let ne0 = params.ne0;
84
85    let i3 = i / (ne2 * ne1 * ne0);
86    i = i % (ne2 * ne1 * ne0);
87    let i2 = i / (ne1 * ne0);
88    i = i % (ne1 * ne0);
89    let i1 = i / ne0;
90    let i0 = i % ne0;
91
92    let src_idx = i0 * params.stride_src0 + i1 * params.stride_src1 + i2 * params.stride_src2 + i3 * params.stride_src3;
93
94#ifdef ABS
95    let res = abs(src[params.offset_src + src_idx]);
96#endif
97#ifdef SGN
98    let res = select(TYPE(select(0.0, -1.0, src[params.offset_src + src_idx] < 0.0)), TYPE(1.0), src[params.offset_src + src_idx] > 0.0);
99#endif
100#ifdef NEG
101    let res = -src[params.offset_src + src_idx];
102#endif
103#ifdef STEP
104    let res = TYPE(select(0.0, 1.0, src[params.offset_src + src_idx] > 0.0));
105#endif
106#ifdef TANH
107    let res = tanh(clamp(src[params.offset_src + src_idx], -9.010913, 9.010913));
108#endif
109#ifdef RELU
110    let res = select(0.0, src[params.offset_src + src_idx], src[params.offset_src + src_idx] > 0.0);
111#endif
112#ifdef ELU
113    let res = select(exp(src[params.offset_src + src_idx]) - 1.0, src[params.offset_src + src_idx], src[params.offset_src + src_idx] > 0.0);
114#endif
115#ifdef HARDSIGMOID
116    let res = min(1.0, max(0.0, (src[params.offset_src + src_idx] + 3.0) / 6.0));
117#endif
118#ifdef SIGMOID
119    let res = 1.0 / (1.0 + exp(-src[params.offset_src + src_idx]));
120#endif
121#ifdef SILU
122    let res = src[params.offset_src + src_idx] / (1.0 + exp(-src[params.offset_src + src_idx]));
123#endif
124#ifdef EXP
125    let src_f32 = f32(src[params.offset_src + src_idx]);
126    let res = TYPE(exp(src_f32));
127#endif
128#ifdef LOG
129    let res = TYPE(log(f32(src[params.offset_src + src_idx])));
130#endif
131#ifdef CLAMP
132    let res = clamp(src[params.offset_src + src_idx], TYPE(params.clamp_min), TYPE(params.clamp_max));
133#endif
134#ifdef FILL
135    let res = TYPE(params.fill_val);
136#endif
137#ifdef HARDSWISH
138    let res = src[params.offset_src + src_idx] * min(1.0, max(0.0, (src[params.offset_src + src_idx] + 3.0) / 6.0));
139#endif
140#ifdef GELU
141    let res = 0.5 * src[params.offset_src + src_idx] * (1.0 + tanh(clamp(0.7978845608028654 * (src[params.offset_src + src_idx] + 0.044715 * src[params.offset_src + src_idx] * src[params.offset_src + src_idx] * src[params.offset_src + src_idx]), -9.010913, 9.010913)));
142#endif
143#ifdef GELU_QUICK
144    let res = src[params.offset_src + src_idx] * (1.0 / (1.0 + exp(clamp(-1.702 * src[params.offset_src + src_idx], -80.0, 80.0))));
145#endif
146#ifdef GELU_ERF
147    let res = 0.5 * src[params.offset_src + src_idx] * (1.0 + erf_approx(src[params.offset_src + src_idx] * 0.7071067811865476));
148#endif
149#ifdef XIELU
150    let val = f32(src[params.offset_src + src_idx]);
151    let res =
152        TYPE(select(
153            ((exp(min(val, params.eps)) - 1.0) - val) * params.alpha_n + params.beta * val,
154            params.alpha_p * val * val + params.beta * val,
155            val > 0.0));
156#endif
157#ifdef SOFTPLUS
158    let src_f32 = f32(src[params.offset_src + src_idx]);
159    let res = TYPE(select(log(1.0 + exp(src_f32)), src_f32, src_f32 > 20.0));
160#endif
161#ifdef EXPM1
162    let src_f32 = f32(src[params.offset_src + src_idx]);
163    let res = TYPE(exp(src_f32) - 1.0);
164#endif
165#ifdef FLOOR
166    let res = floor(src[params.offset_src + src_idx]);
167#endif
168#ifdef CEIL
169    let res = ceil(src[params.offset_src + src_idx]);
170#endif
171#ifdef ROUND
172    let src_f32 = f32(src[params.offset_src + src_idx]);
173    let result = select(ceil(src_f32 - 0.5), floor(src_f32 + 0.5), src_f32 >= 0.0);
174    let res = TYPE(result);
175#endif
176#ifdef TRUNC
177    let res = trunc(src[params.offset_src + src_idx]);
178#endif
179#ifdef SQR
180    let res = src[params.offset_src + src_idx] * src[params.offset_src + src_idx];
181#endif
182#ifdef SQRT
183    let res = TYPE(sqrt(f32(src[params.offset_src + src_idx])));
184#endif
185#ifdef SIN
186    let res_f32 = sin(f32(src[params.offset_src + src_idx]));
187    let res = TYPE(res_f32);
188#endif
189#ifdef COS
190    let res_f32 = cos(f32(src[params.offset_src + src_idx]));
191    let res = TYPE(res_f32);
192#endif
193#ifdef DIAG
194    let res = select(0.0, src[params.offset_src + i0 + i2 * params.stride_src2 + i3 * params.stride_src3], i0 == i1);
195#endif
196#ifdef TRI
197#ifdef TRI_TYPE_LOWER
198    let res = select(0.0, src[params.offset_src + src_idx], i0 < i1);
199#elif TRI_TYPE_LOWER_DIAG
200    let res = select(0.0, src[params.offset_src + src_idx], i0 <= i1);
201#elif TRI_TYPE_UPPER
202    let res = select(0.0, src[params.offset_src + src_idx], i0 > i1);
203#elif TRI_TYPE_UPPER_DIAG
204    let res = select(0.0, src[params.offset_src + src_idx], i0 >= i1);
205#endif
206#endif
207
208#ifdef INPLACE
209    src[params.offset_src + src_idx] = res;
210#else
211    dst[params.offset_dst + flat_i] = res;
212#endif
213}