char-fork/llama.cpp

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

shaofeiqiopencl: add cumsum op (#18981)3d9ab225e

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.8 KiB139 linesraw
1#pragma OPENCL EXTENSION cl_khr_fp16 : enable
2
3#ifdef cl_intel_required_subgroup_size
4#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
5#define INTEL_GPU 1
6#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
7#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
8#elif defined(cl_qcom_reqd_sub_group_size)
9#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
10#define ADRENO_GPU 1
11#define REQD_SUBGROUP_SIZE_64  __attribute__((qcom_reqd_sub_group_size("half")))
12#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
13#endif
14
15// max workgroup size is usually 1024, this covers various subgroups sizes
16#define MAX_SUBGROUPS 128
17
18#ifdef INTEL_GPU
19REQD_SUBGROUP_SIZE_32
20#elif defined (ADRENO_GPU)
21REQD_SUBGROUP_SIZE_64
22#endif
23kernel void kernel_cumsum_blk(
24        global char * src0,
25        ulong offset0,
26        global char * tmp,
27        global char * dst,
28        ulong offsetd,
29        int   ne00,
30        int   ne01,
31        int   ne02,
32        int   ne03,
33        ulong nb00,
34        ulong nb01,
35        ulong nb02,
36        ulong nb03,
37        uint net0,
38        uint net1,
39        uint net2
40) {
41    src0 = src0 + offset0;
42    dst  = dst + offsetd;
43
44    const int i3 = get_group_id(2);
45    const int i2 = get_group_id(1);
46    const int i1 = get_group_id(0);
47
48    const int nth = get_local_size(0);
49    const int tid = get_local_id(0);
50
51    const uint sg_size = get_sub_group_size();
52    const uint sg_id = get_sub_group_id();
53    const uint sg_lid = get_sub_group_local_id();
54
55    const int ib = i1 / ne01;
56    const int i00 = ib * nth;
57    const int i01 = i1 % ne01;
58    const int i02 = i2;
59    const int i03 = i3;
60
61    global const float * src0_row = (global const float *)(src0 + i03*nb03 + i02*nb02 + i01*nb01);
62    global       float * tmp_row  = (global float *)tmp + net0 * i01 + net0 * net1 * i02 + net0 * net1 * net2 * i03;
63    global       float * dst_row  = (global float *)dst + i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00;
64
65    __local float partial[MAX_SUBGROUPS];
66
67    float v = 0.0f;
68    if (i00 + tid < ne00) {
69        v = src0_row[i00 + tid];
70    }
71
72    float s = sub_group_scan_inclusive_add(v);
73    if (sg_lid == sg_size - 1) {
74        partial[sg_id] = s;
75    }
76    barrier(CLK_LOCAL_MEM_FENCE);
77
78    // NB: subgroup size should be larger than number of subgroups
79    // assuming max workgroup size of 1024, subgroup size should be >= 32
80    if (sg_id == 0) {
81        float x = 0.0f;
82        if (sg_lid < get_num_sub_groups()) {
83            x = partial[sg_lid];
84        }
85        float ex = sub_group_scan_exclusive_add(x);
86        if (sg_lid < get_num_sub_groups()) {
87            partial[sg_lid] = ex;
88        }
89    }
90    barrier(CLK_LOCAL_MEM_FENCE);
91
92    s += partial[sg_id];
93
94    if (i00 + tid < ne00) {
95        dst_row[i00 + tid] = s;
96    }
97    if (ne00 > nth && tid == nth - 1) {
98        tmp_row[ib] = s;
99    }
100}
101
102kernel void kernel_cumsum_add(
103        global char * tmp,
104        global char * dst,
105        ulong offsetd,
106        int   ne00,
107        int   ne01,
108        int   ne02,
109        int   ne03,
110        uint nbt0,
111        uint nbt1,
112        uint nbt2,
113        uint nbt3
114) {
115    dst  = dst + offsetd;
116
117    const int i3 = get_group_id(2);
118    const int i2 = get_group_id(1);
119    const int i1 = get_group_id(0);
120
121    const int nth = get_local_size(0);
122    const int tid = get_local_id(0);
123
124    const int ib = i1 / ne01;
125    if (ib == 0) {
126        return;
127    }
128    const int i00 = ib * nth;
129    const int i01 = i1 % ne01;
130    const int i02 = i2;
131    const int i03 = i3;
132
133    global float * tmp_row  = (global float *)(tmp + nbt1 * i01 + nbt2 * i02 + nbt3 * i03);
134    global float * dst_row  = (global float *)dst + i03*ne02*ne01*ne00 + i02*ne01*ne00 + i01*ne00;
135
136    if (i00 + tid < ne00) {
137        dst_row[i00 + tid] += tmp_row[ib - 1];
138    }
139}