char-fork/llama.cpp

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

Chen Yuanggml-webgpu: add the upscale shader (#22419)aab68217b

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
7.3 KiB240 linesraw
1#if defined(SRC_F16) || defined(DST_F16)
2enable f16;
3#endif
4
5#ifdef SRC_F16
6#define SRC_TYPE f16
7#else
8#define SRC_TYPE f32
9#endif
10
11#ifdef DST_F16
12#define DST_TYPE f16
13#else
14#define DST_TYPE f32
15#endif
16
17@group(0) @binding(0)
18var<storage, read_write> input: array<SRC_TYPE>;
19
20@group(0) @binding(1)
21var<storage, read_write> output: array<DST_TYPE>;
22
23struct Params {
24    offset_i: u32,
25    offset_o: u32,
26
27    // element strides
28    si0: u32, si1: u32, si2: u32, si3: u32,
29    so0: u32, so1: u32, so2: u32, so3: u32,
30
31    src_w: u32,
32    src_h: u32,
33    src_z: u32,
34    src_n: u32,
35
36    dst_w: u32,
37    dst_h: u32,
38    dst_z: u32,
39    dst_n: u32,
40
41    mode_flags: u32,
42};
43
44@group(0) @binding(2)
45var<uniform> params: Params;
46
47const GGML_SCALE_FLAG_ALIGN_CORNERS: u32 = 1u << 8u;
48
49fn get_clamped_input(x: i32, y: i32, z: u32, n: u32) -> f32 {
50    let cx = u32(clamp(x, 0, i32(params.src_w) - 1));
51    let cy = u32(clamp(y, 0, i32(params.src_h) - 1));
52    let i = params.offset_i + cx * params.si0 + cy * params.si1 + z * params.si2 + n * params.si3;
53    return f32(input[i]);
54}
55
56fn cubic_weight(t: f32, a: f32) -> f32 {
57    let at = abs(t);
58    if (at <= 1.0) {
59        return (a + 2.0) * at * at * at - (a + 3.0) * at * at + 1.0;
60    } else if (at <= 2.0) {
61        return a * at * at * at - 5.0 * a * at * at + 8.0 * a * at - 4.0 * a;
62    } else {
63        return 0.0;
64    }
65}
66
67@compute @workgroup_size(WG_SIZE)
68fn main(
69    @builtin(global_invocation_id) gid: vec3<u32>,
70    @builtin(num_workgroups) num_wg: vec3<u32>
71) {
72
73    let i_out = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y;
74    let total = params.dst_w * params.dst_h * params.dst_z * params.dst_n;
75
76    if (i_out >= total) {
77        return;
78    }
79
80    // decode (x, y, z, n)
81    var i = i_out;
82    let x_dst = i % params.dst_w;
83    i = i / params.dst_w;
84    let y_dst = i % params.dst_h;
85    i = i / params.dst_h;
86    let z_dst = i % params.dst_z;
87    let n_dst = i / params.dst_z;
88
89    // scale factors
90    var sf0 = f32(params.dst_w) / f32(params.src_w);
91    var sf1 = f32(params.dst_h) / f32(params.src_h);
92    var sf2 = f32(params.dst_z) / f32(params.src_z);
93    var sf3 = f32(params.dst_n) / f32(params.src_n);
94
95    let align_corners = (params.mode_flags & GGML_SCALE_FLAG_ALIGN_CORNERS) != 0;
96
97    // pixel_offset: 0.5 for half-pixel-center (default), 0.0 for align_corners
98    var pixel_offset = 0.5;
99    if (align_corners) {
100        pixel_offset = 0.0;
101        if (params.dst_w > 1 && params.src_w > 1) {
102            sf0 = f32(params.dst_w - 1) / f32(params.src_w - 1);
103        }
104        if (params.dst_h > 1 && params.src_h > 1) {
105            sf1 = f32(params.dst_h - 1) / f32(params.src_h - 1);
106        }
107    }
108
109    let z_src = min(params.src_z - 1, u32(floor(f32(z_dst) / sf2)));
110    let n_src = min(params.src_n - 1, u32(floor(f32(n_dst) / sf3)));
111
112    var result = 0.0;
113
114#if defined(NEAREST)
115
116    let x_src = min(params.src_w - 1, u32(floor(f32(x_dst) / sf0)));
117    let y_src = min(params.src_h - 1, u32(floor(f32(y_dst) / sf1)));
118
119    result = get_clamped_input(i32(x_src), i32(y_src), z_src, n_src);
120
121#elif defined(BILINEAR)
122
123#if defined(ANTIALIAS)
124
125    // Antialiased bilinear: triangle filter over a variable support region.
126    let support0 = max(1.0f / sf0, 1.0f);
127    let support1 = max(1.0f / sf1, 1.0f);
128    let invscale0 = 1.0 / support0;
129    let invscale1 = 1.0 / support1;
130
131    let fx = (f32(x_dst) + pixel_offset) / sf0;
132    let fy = (f32(y_dst) + pixel_offset) / sf1;
133
134    let x_min = max(i32(fx - support0 + pixel_offset), 0);
135    let y_min = max(i32(fy - support1 + pixel_offset), 0);
136    let x_max = min(i32(fx + support0 + pixel_offset), i32(params.src_w));
137    let y_max = min(i32(fy + support1 + pixel_offset), i32(params.src_h));
138
139    var weighted_sum = 0.0;
140    var total_weight = 0.0;
141
142    for (var x = x_min; x < x_max; x += 1) {
143        let wx = max(1.0 - abs(f32(x) - fx + pixel_offset) * invscale0, 0.0);
144        for (var y = y_min; y < y_max; y += 1) {
145            let wy = max(1.0 - abs(f32(y) - fy + pixel_offset) * invscale1, 0.0);
146            let w = wx * wy;
147            if (w > 0.0) {
148                weighted_sum += get_clamped_input(x, y, z_src, n_src) * w;
149                total_weight += w;
150            }
151        }
152    }
153
154    if (total_weight > 0.0) {
155        result = weighted_sum / total_weight;
156    }
157
158#else
159
160    let fx = (f32(x_dst) + pixel_offset) / sf0 - pixel_offset;
161    let fy = (f32(y_dst) + pixel_offset) / sf1 - pixel_offset;
162    let x0 = i32(floor(fx));
163    let y0 = i32(floor(fy));
164    let dx = clamp(fx - f32(x0), 0.0, 1.0);
165    let dy = clamp(fy - f32(y0), 0.0, 1.0);
166    let a = get_clamped_input(x0, y0, z_src, n_src);
167    let b = get_clamped_input(x0 + 1, y0, z_src, n_src);
168    let c = get_clamped_input(x0, y0 + 1, z_src, n_src);
169    let d = get_clamped_input(x0 + 1, y0 + 1, z_src, n_src);
170
171    let wa = (1.0 - dx) * (1.0 - dy);
172    let wb = dx * (1.0 - dy);
173    let wc = (1.0 - dx) * dy;
174    let wd = dx * dy;
175
176    result = a * wa + b * wb + c * wc + d * wd;
177
178#endif
179
180#elif defined(BICUBIC)
181
182    // bicubic convolution with alpha = -0.75 (PyTorch default)
183    let alpha = -0.75;
184    let fx = (f32(x_dst) + pixel_offset) / sf0 - pixel_offset;
185    let fy = (f32(y_dst) + pixel_offset) / sf1 - pixel_offset;
186
187    let x0 = i32(floor(fx));
188    let y0 = i32(floor(fy));
189    let dx = fx - f32(x0);
190    let dy = fy - f32(y0);
191
192    // horizontal weights for offsets -1, 0, 1, 2
193    let wx0 = cubic_weight(dx + 1.0, alpha);
194    let wx1 = cubic_weight(dx, alpha);
195    let wx2 = cubic_weight(1.0 - dx, alpha);
196    let wx3 = cubic_weight(2.0 - dx, alpha);
197
198    // vertical weights for offsets -1, 0, 1, 2
199    let wy0 = cubic_weight(dy + 1.0, alpha);
200    let wy1 = cubic_weight(dy, alpha);
201    let wy2 = cubic_weight(1.0 - dy, alpha);
202    let wy3 = cubic_weight(2.0 - dy, alpha);
203
204    // intermediate horizontal interpolation for 4x4 grid of pixels
205    // x0-1, x0, x0+1, x0+2, y0-1
206    let p0 = get_clamped_input(x0 - 1, y0 - 1, z_src, n_src);
207    let p1 = get_clamped_input(x0, y0 - 1, z_src, n_src);
208    let p2 = get_clamped_input(x0 + 1, y0 - 1, z_src, n_src);
209    let p3 = get_clamped_input(x0 + 2, y0 - 1, z_src, n_src);
210    let row0 = p0 * wx0 + p1 * wx1 + p2 * wx2 + p3 * wx3;
211
212    // x0-1, x0, x0+1, x0+2, y0
213    let q0 = get_clamped_input(x0 - 1, y0, z_src, n_src);
214    let q1 = get_clamped_input(x0, y0, z_src, n_src);
215    let q2 = get_clamped_input(x0 + 1, y0, z_src, n_src);
216    let q3 = get_clamped_input(x0 + 2, y0, z_src, n_src);
217    let row1 = q0 * wx0 + q1 * wx1 + q2 * wx2 + q3 * wx3;
218
219    // x0-1, x0, x0+1, x0+2, y0+1
220    let r0 = get_clamped_input(x0 - 1, y0 + 1, z_src, n_src);
221    let r1 = get_clamped_input(x0, y0 + 1, z_src, n_src);
222    let r2 = get_clamped_input(x0 + 1, y0 + 1, z_src, n_src);
223    let r3 = get_clamped_input(x0 + 2, y0 + 1, z_src, n_src);
224    let row2 = r0 * wx0 + r1 * wx1 + r2 * wx2 + r3 * wx3;
225
226    // x0-1, x0, x0+1, x0+2, y0+2
227    let s0 = get_clamped_input(x0 - 1, y0 + 2, z_src, n_src);
228    let s1 = get_clamped_input(x0, y0 + 2, z_src, n_src);
229    let s2 = get_clamped_input(x0 + 1, y0 + 2, z_src, n_src);
230    let s3 = get_clamped_input(x0 + 2, y0 + 2, z_src, n_src);
231    let row3 = s0 * wx0 + s1 * wx1 + s2 * wx2 + s3 * wx3;
232
233    // final vertical interpolation
234    result = row0 * wy0 + row1 * wy1 + row2 * wy2 + row3 * wy3;
235
236#endif
237
238    let dst_idx = params.offset_o + x_dst * params.so0 + y_dst * params.so1 + z_dst * params.so2 + n_dst * params.so3;
239    output[dst_idx] = DST_TYPE(result);
240}