char-fork/llama.cpp

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

Masashi Yoshimuraggml-webgpu : refactor several wgsl files and simplify flash_attn wgsl. (#26134)f401bb139

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.6 KiB100 linesraw
1#include "common_decls.tmpl"
2enable f16;
3
4// Ported from the Vulkan backend's conv2d_dw.comp. Two variants (based on WHCN)
5// selected by the input (src1) layout: contiguous -> WHCN, else CWHN.
6// weight (src0) is [KW,KH,1,C]; output matches the input layout.
7
8@group(0) @binding(0)
9var<storage, read_write> weights: array<WEIGHT_TYPE>;
10@group(0) @binding(1)
11var<storage, read_write> input: array<INPUT_TYPE>;
12@group(0) @binding(2)
13var<storage, read_write> output: array<OUTPUT_TYPE>;
14
15struct Params {
16    offset_w: u32,
17    offset_i: u32,
18    offset_o: u32,
19
20    ne: u32,
21    channels: u32,
22    dst_w: u32, dst_h: u32,
23    src_w: u32, src_h: u32,
24    knl_w: u32, knl_h: u32,
25
26    stride_x: i32, stride_y: i32,
27    pad_x: i32, pad_y: i32,
28    dilation_x: i32, dilation_y: i32,
29};
30
31@group(0) @binding(3)
32var<uniform> params: Params;
33
34#if defined(WHCN)
35// Input/output/kernel contiguous in [W, H, C, N] order (kernel [KW,KH,C]).
36fn conv_2d_dw(idx: u32) -> f32 {
37    let i0    = idx / params.dst_w;
38    let dst_x = idx - i0 * params.dst_w;
39    let i1    = i0 / params.dst_h;
40    let dst_y = i0 - i1 * params.dst_h;
41    let n     = i1 / params.channels;
42    let c     = i1 - n * params.channels;
43
44    let src_i = params.offset_i + n * params.channels * params.src_h * params.src_w
45                                + c * params.src_h * params.src_w;
46    let knl_i = params.offset_w + c * params.knl_h * params.knl_w;
47
48    var sum: f32 = 0.0;
49    for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) {
50        let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y;
51        if (src_y < 0 || src_y >= i32(params.src_h)) { continue; }
52        for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
53            let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
54            if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
55            let v = f32(input[src_i + u32(src_y) * params.src_w + u32(src_x)]);
56            let k = f32(weights[knl_i + ky * params.knl_w + kx]);
57            sum += v * k;
58        }
59    }
60    return sum;
61}
62#else
63// Channels contiguous (CWHN): channel is the innermost axis.
64fn conv_2d_dw(idx: u32) -> f32 {
65    let i0    = idx / params.channels;
66    let c     = idx - i0 * params.channels;
67    let i1    = i0 / params.dst_w;
68    let dst_x = i0 - i1 * params.dst_w;
69    let n     = i1 / params.dst_h;
70    let dst_y = i1 - n * params.dst_h;
71
72    let src_i   = params.offset_i + n * params.channels * params.src_h * params.src_w;
73    let src_row = params.src_w * params.channels;
74    let knl_row = params.knl_w * params.channels;
75
76    var sum: f32 = 0.0;
77    for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) {
78        let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y;
79        if (src_y < 0 || src_y >= i32(params.src_h)) { continue; }
80        for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
81            let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
82            if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
83            let v = f32(input[src_i + u32(src_y) * src_row + u32(src_x) * params.channels + c]);
84            let k = f32(weights[params.offset_w + ky * knl_row + kx * params.channels + c]);
85            sum += v * k;
86        }
87    }
88    return sum;
89}
90#endif
91
92@compute @workgroup_size(WG_SIZE)
93fn main(
94    @builtin(global_invocation_id) gid: vec3<u32>,
95    @builtin(num_workgroups) num_wg: vec3<u32>
96) {
97    let idx = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y;
98    if (idx >= params.ne) { return; }
99    output[params.offset_o + idx] = OUTPUT_TYPE(conv_2d_dw(idx));
100}