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 →
3.4 KiB142 linesraw
1enable f16;
2
3struct Params {
4    ne: u32,
5
6    // offsets in elements
7    offset_src0: u32,
8    offset_src1: u32,
9    offset_dst: u32,
10
11    stride_src0_0: u32,
12    stride_src0_1: u32,
13    stride_src0_2: u32,
14    stride_src0_3: u32,
15
16    stride_src1_0: u32,
17    stride_src1_1: u32,
18    stride_src1_2: u32,
19    stride_src1_3: u32,
20
21    a_ne0: u32,
22    a_ne1: u32,
23    a_ne2: u32,
24
25    b_ne0: u32,
26    b_ne1: u32,
27    b_ne2: u32,
28    b_ne3: u32,
29};
30
31fn src0_index(_i: u32) -> u32 {
32    var i = _i;
33    let a_i3 = i / (params.a_ne2 * params.a_ne1 * params.a_ne0);
34    i = i % (params.a_ne2 * params.a_ne1 * params.a_ne0);
35    let a_i2 = i / (params.a_ne1 * params.a_ne0);
36    i = i % (params.a_ne1 * params.a_ne0);
37    let a_i1 = i / params.a_ne0;
38    let a_i0 = i % params.a_ne0;
39
40    return a_i0 * params.stride_src0_0 +
41           a_i1 * params.stride_src0_1 +
42           a_i2 * params.stride_src0_2 +
43           a_i3 * params.stride_src0_3;
44}
45
46fn src1_index(_i: u32) -> u32 {
47    var i = _i;
48    let a_i3 = i / (params.a_ne2 * params.a_ne1 * params.a_ne0);
49    i = i % (params.a_ne2 * params.a_ne1 * params.a_ne0);
50    let a_i2 = i / (params.a_ne1 * params.a_ne0);
51    i = i % (params.a_ne1 * params.a_ne0);
52    let a_i1 = i / params.a_ne0;
53    let a_i0 = i % params.a_ne0;
54
55    // handle repetition of b
56    // index loops back to the beginning and repeats after elements are exhausted = modulo
57    let b_i0 = a_i0 % params.b_ne0;
58    let b_i1 = a_i1 % params.b_ne1;
59    let b_i2 = a_i2 % params.b_ne2;
60    let b_i3 = a_i3 % params.b_ne3;
61
62    // compute index for position in b's flat array
63    return b_i0 * params.stride_src1_0 +
64           b_i1 * params.stride_src1_1 +
65           b_i2 * params.stride_src1_2 +
66           b_i3 * params.stride_src1_3;
67}
68
69#ifdef TYPE_F32
70#define DataType f32
71#endif
72#ifdef TYPE_F16
73#define DataType f16
74#endif
75
76#ifdef SRC_OVERLAP
77@group(0) @binding(0)
78var<storage, read_write> merged_src: array<DataType>;
79
80@group(0) @binding(1)
81var<storage, read_write> dst: array<DataType>;
82
83@group(0) @binding(2)
84var<uniform> params: Params;
85#else
86@group(0) @binding(0)
87var<storage, read_write> src0: array<DataType>;
88
89@group(0) @binding(1)
90var<storage, read_write> src1 : array<DataType>;
91#if defined(INPLACE) || defined(OVERLAP)
92@group(0) @binding(2)
93var<uniform> params: Params;
94
95#else
96@group(0) @binding(2)
97var<storage, read_write> dst: array<DataType>;
98
99@group(0) @binding(3)
100var<uniform> params: Params;
101#endif
102#endif
103
104fn op(a: DataType, b: DataType) -> DataType {
105#ifdef OP_ADD
106    return a + b;
107#elif defined(OP_SUB)
108    return a - b;
109#elif defined(OP_MUL)
110    return a * b;
111#elif defined(OP_DIV)
112    return a / b;
113#endif
114}
115
116fn update(dst_i: u32, src0_i: u32, src1_i: u32) {
117#ifdef SRC_OVERLAP
118    let result = op(merged_src[src0_i], merged_src[src1_i]);
119#else
120    let result = op(src0[src0_i], src1[src1_i]);
121#endif
122
123#ifdef INPLACE
124    src0[src0_i] = result;
125#elif defined(OVERLAP)
126    src1[src1_i] = result;
127#else
128    dst[dst_i] = result;
129#endif
130}
131
132@compute @workgroup_size(WG_SIZE)
133fn main(@builtin(global_invocation_id) gid: vec3<u32>,
134    @builtin(num_workgroups)       num_wg:  vec3<u32>) {
135    let threads_per_group = u32(WG_SIZE);
136    let i = gid.x + (num_wg.x * threads_per_group) * gid.y;
137    if (i < params.ne) {
138        let src0_i = params.offset_src0 + src0_index(i);
139        let src1_i = params.offset_src1 + src1_index(i);
140        update(params.offset_dst + i, src0_i, src1_i);
141    }
142}