char-fork/llama.cpp

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

shaofeiqiopencl: optimize mean and sum_row kernels (#19614)983559d24

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.3 KiB140 linesraw
1#pragma OPENCL EXTENSION cl_khr_fp16 : enable
2#pragma OPENCL EXTENSION cl_khr_subgroups : enable
3
4// Most devices have max workgroup size of 1024, so this is enough for subgroup
5// sizes of 16, 32, 64 and 128. Increase this value for smaller subgroups sizes
6#define MAX_SUBGROUPS 64
7kernel void kernel_mean_f32(
8    global char *  src0,
9    ulong           offset0,
10    global char *  dst,
11    ulong           offsetd,
12    int             ne00,
13    int             ne01,
14    int             ne02,
15    int             ne03,
16    ulong           nb01,
17    ulong           nb02,
18    ulong           nb03,
19    ulong           nb1,
20    ulong           nb2,
21    ulong           nb3
22) {
23    src0 = src0 + offset0;
24    dst  = dst  + offsetd;
25
26    const int i3 = get_group_id(2);
27    const int i2 = get_group_id(1);
28    const int i1 = get_group_id(0);
29
30    const int lid = get_local_id(0);
31    const int lsize = get_local_size(0);
32
33    const uint sg_size = get_sub_group_size();
34    const uint sg_id = get_sub_group_id();
35    const uint sg_lid = get_sub_group_local_id();
36
37    __local float lmem[MAX_SUBGROUPS];
38
39    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
40        return;
41    }
42
43    if(sg_id == 0){
44        lmem[sg_lid] = 0.0f;
45    }
46
47    global float * src_row = (global float *) (src0 + i1*nb01 + i2*nb02 + i3*nb03);
48    global float * dst_row = (global float *) (dst  + i1*nb1  + i2*nb2  + i3*nb3);
49
50    float sumf = 0.0f;
51
52    for (int i0 = lid; i0 < ne00; i0 += lsize) {
53        sumf += src_row[i0];
54    }
55
56    sumf = sub_group_reduce_add(sumf);
57
58    barrier(CLK_LOCAL_MEM_FENCE);
59
60    if(sg_lid == 0){
61        lmem[sg_id] = sumf;
62    }
63
64    barrier(CLK_LOCAL_MEM_FENCE);
65
66    sumf = lmem[sg_lid];
67    sumf = sub_group_reduce_add(sumf);
68
69    if (lid == 0) {
70        dst_row[0] = sumf / ne00;
71    }
72}
73
74kernel void kernel_mean_f32_4(
75    global char *  src0,
76    ulong           offset0,
77    global char *  dst,
78    ulong           offsetd,
79    int             ne00,
80    int             ne01,
81    int             ne02,
82    int             ne03,
83    ulong           nb01,
84    ulong           nb02,
85    ulong           nb03,
86    ulong           nb1,
87    ulong           nb2,
88    ulong           nb3
89) {
90    src0 = src0 + offset0;
91    dst  = dst  + offsetd;
92
93    const int i3 = get_group_id(2);
94    const int i2 = get_group_id(1);
95    const int i1 = get_group_id(0);
96
97    const int lid = get_local_id(0);
98    const int lsize = get_local_size(0);
99
100    const uint sg_size = get_sub_group_size();
101    const uint sg_id = get_sub_group_id();
102    const uint sg_lid = get_sub_group_local_id();
103
104    __local float lmem[MAX_SUBGROUPS];
105
106    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
107        return;
108    }
109
110    if(sg_id == 0){
111        lmem[sg_lid] = 0.0f;
112    }
113
114    global float4 * src_row = (global float4 *) (src0 + i1*nb01 + i2*nb02 + i3*nb03);
115    global float  * dst_row = (global float  *) (dst  + i1*nb1  + i2*nb2  + i3*nb3);
116
117    float4 sum_vec = (float4)0.0f;
118
119    for (int i0 = lid; i0 < ne00 / 4; i0 += lsize) {
120        sum_vec += src_row[i0];
121    }
122
123    float sumf = dot(sum_vec, (float4)(1.0f));
124    sumf = sub_group_reduce_add(sumf);
125
126    barrier(CLK_LOCAL_MEM_FENCE);
127
128    if(sg_lid == 0){
129        lmem[sg_id] = sumf;
130    }
131
132    barrier(CLK_LOCAL_MEM_FENCE);
133
134    sumf = lmem[sg_lid];
135    sumf = sub_group_reduce_add(sumf);
136
137    if (lid == 0) {
138        dst_row[0] = sumf / ne00;
139    }
140}