Skip to content

Commit 5157172

Browse files
lhezlhez
authored
opencl: optimize mul_mat_f16_f32_l4 for decode (#24504)
1 parent cda6385 commit 5157172

2 files changed

Lines changed: 348 additions & 6 deletions

File tree

‎ggml/src/ggml-opencl/ggml-opencl.cpp‎

Lines changed: 52 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -564,6 +564,9 @@ struct ggml_backend_opencl_context {
564564
cl_kernel kernel_mul_mat_f16_f32_1row;
565565
cl_kernel kernel_mul_mat_f16_f32;
566566
cl_kernel kernel_mul_mat_f16_f32_l4;
567+
cl_kernel kernel_mul_mat_f16_f32_l4_dr;
568+
cl_kernel kernel_mul_mat_f16_f32_l4_dr_ls;
569+
cl_kernel kernel_mul_mat_f16_f32_l4_dr_lq;
567570
cl_kernel kernel_mul_mat_f16_f32_tiled;
568571
cl_kernel kernel_adreno_xmem_pack_src_f32;
569572
cl_kernel kernel_adreno_xmem_prepack_weight_f16;
@@ -1787,6 +1790,11 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
17871790
build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src.c_str(), compile_opts);
17881791

17891792
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_l4 = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4", &err), err));
1793+
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_l4_dr = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_dr", &err), err));
1794+
if (backend_ctx->gpu_family == ADRENO) {
1795+
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_l4_dr_ls = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_dr_ls", &err), err));
1796+
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_l4_dr_lq = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_dr_lq", &err), err));
1797+
}
17901798
GGML_LOG_CONT(".");
17911799
}
17921800

@@ -14570,11 +14578,31 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
1457014578
}
1457114579

1457214580
if (src1t == GGML_TYPE_F32) {
14581+
// heuristic for packing more work for Adreno
14582+
const bool adreno_use_lane_split =
14583+
backend_ctx->gpu_family == ADRENO &&
14584+
ne11 == 1 &&
14585+
ne01 >= 8 &&
14586+
ne00 % 4 == 0 &&
14587+
r3 == 1 && r2 >= 1 && r2 <= 8 &&
14588+
(ne12 % r2) == 0;
14589+
1457314590
if (ne11 * ne12 < 4) {
1457414591
kernel = backend_ctx->kernel_mul_mat_f16_f32_1row;
14592+
} else if (adreno_use_lane_split && ne00 >= 64 && ne00 <= 128) {
14593+
kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_dr_lq;
14594+
nrows = 1;
14595+
} else if (adreno_use_lane_split && r2 >= 2 && ne00 > 128 && ne00 <= 256) {
14596+
kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_dr_ls;
14597+
nrows = 1;
1457514598
} else if (ne00 >= 128 && ne01 >= 8 && ne00%4 == 0) {
14576-
kernel = backend_ctx->kernel_mul_mat_f16_f32_l4;
14577-
nrows = ne11;
14599+
if (ne11 == 1) {
14600+
kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_dr;
14601+
nrows = 1; // not used by this kernel
14602+
} else {
14603+
kernel = backend_ctx->kernel_mul_mat_f16_f32_l4;
14604+
nrows = ne11;
14605+
}
1457814606
} else {
1457914607
kernel = backend_ctx->kernel_mul_mat_f16_f32;
1458014608
nrows = 4;
@@ -15353,12 +15381,30 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co
1535315381

1535415382
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
1535515383
} else {
15356-
int64_t ny = (ne11 + nrows - 1)/nrows;
15384+
if (kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_dr) {
15385+
const int NDST_DR = 4;
15386+
size_t global_work_size[] = {(size_t)CEIL_DIV(ne01, NDST_DR)*nth0, (size_t)nth1, (size_t)ne12*ne13};
15387+
size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
1535715388

15358-
size_t global_work_size[] = {(size_t)ne01*nth0, (size_t)ny*nth1, (size_t)ne12*ne13};
15359-
size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
15389+
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
15390+
} else if (kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_dr_ls) {
15391+
size_t global_work_size[] = {(size_t)CEIL_DIV(ne01, 2)*nth0, (size_t)nth1, (size_t)ne02*ne03};
15392+
size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
1536015393

15361-
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
15394+
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
15395+
} else if (kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_dr_lq) {
15396+
size_t global_work_size[] = {(size_t)CEIL_DIV(ne01, 4)*nth0, (size_t)nth1, (size_t)ne02*ne03};
15397+
size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
15398+
15399+
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
15400+
} else {
15401+
int64_t ny = (ne11 + nrows - 1)/nrows;
15402+
15403+
size_t global_work_size[] = {(size_t)ne01*nth0, (size_t)ny*nth1, (size_t)ne12*ne13};
15404+
size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
15405+
15406+
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
15407+
}
1536215408
}
1536315409
}
1536415410

‎ggml/src/ggml-opencl/kernels/mul_mv_f16_f32_l4.cl‎

Lines changed: 296 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,3 +82,299 @@ kernel void kernel_mul_mat_f16_f32_l4(
8282
}
8383
}
8484
}
85+
86+
// Each subgroup produces DR_NDST outputs, assumes ne11 == 1
87+
#define MUL_MAT_F16_F32_L4_DR_NDST 4
88+
89+
#ifdef ADRENO_GPU
90+
REQD_SUBGROUP_SIZE_64
91+
#endif
92+
kernel void kernel_mul_mat_f16_f32_l4_dr(
93+
global char * src0,
94+
ulong offset0,
95+
global char * src1,
96+
ulong offset1,
97+
global float * dst,
98+
ulong offsetd,
99+
int ne00,
100+
int ne01,
101+
int ne02,
102+
ulong nb00,
103+
ulong nb01,
104+
ulong nb02,
105+
ulong nb03,
106+
int ne10,
107+
int ne11,
108+
int ne12,
109+
ulong nb10,
110+
ulong nb11,
111+
ulong nb12,
112+
ulong nb13,
113+
int ne0,
114+
int ne1,
115+
int r2,
116+
int r3
117+
) {
118+
src0 = (global char*)((global char*)src0 + offset0);
119+
src1 = (global char*)((global char*)src1 + offset1);
120+
dst = (global float*)((global char*)dst + offsetd);
121+
122+
const int r0_base = get_group_id(0) * MUL_MAT_F16_F32_L4_DR_NDST;
123+
const int im = get_group_id(2);
124+
125+
const int i12 = im % ne12;
126+
const int i13 = im / ne12;
127+
128+
// assume ne11 == 1
129+
const ulong offset_src1 = i12*nb12 + i13*nb13;
130+
global float4 * y4 = (global float4 *)(src1 + offset_src1);
131+
132+
global half4 * x4[MUL_MAT_F16_F32_L4_DR_NDST];
133+
float sumf[MUL_MAT_F16_F32_L4_DR_NDST];
134+
135+
const ulong k_head_off = (i12/r2)*nb02 + (i13/r3)*nb03;
136+
137+
#pragma unroll
138+
for (int n = 0; n < MUL_MAT_F16_F32_L4_DR_NDST; ++n) {
139+
int r0 = r0_base + n;
140+
int r0c = r0 < ne01 ? r0 : 0;
141+
ulong off = (ulong)r0c*nb01 + k_head_off;
142+
x4[n] = (global half4 *)(src0 + off);
143+
sumf[n] = 0.0f;
144+
}
145+
146+
const int n_chunks = ne00 / 4;
147+
const int sg_size = get_max_sub_group_size();
148+
const int lid = get_sub_group_local_id();
149+
150+
for (int i = lid; i < n_chunks; i += sg_size) {
151+
float4 q = y4[i];
152+
#pragma unroll
153+
for (int n = 0; n < MUL_MAT_F16_F32_L4_DR_NDST; ++n) {
154+
float4 k = convert_float4(x4[n][i]);
155+
sumf[n] = mad(k.s0, q.s0, sumf[n]);
156+
sumf[n] = mad(k.s1, q.s1, sumf[n]);
157+
sumf[n] = mad(k.s2, q.s2, sumf[n]);
158+
sumf[n] = mad(k.s3, q.s3, sumf[n]);
159+
}
160+
}
161+
162+
#pragma unroll
163+
for (int n = 0; n < MUL_MAT_F16_F32_L4_DR_NDST; ++n) {
164+
float reduced = sub_group_reduce_add(sumf[n]);
165+
int r0 = r0_base + n;
166+
if (lid == 0 && r0 < ne01) {
167+
dst[im*ne1*ne0 + r0] = reduced;
168+
}
169+
}
170+
}
171+
172+
// Kernels for decoding, Adreno only for now
173+
#define MUL_MAT_F16_F32_L4_DR_LS_R2_MAX 8
174+
175+
#ifdef ADRENO_GPU
176+
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
177+
#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
178+
179+
REQD_SUBGROUP_SIZE_64
180+
kernel void kernel_mul_mat_f16_f32_l4_dr_ls(
181+
global char * src0,
182+
ulong offset0,
183+
global char * src1,
184+
ulong offset1,
185+
global float * dst,
186+
ulong offsetd,
187+
int ne00,
188+
int ne01,
189+
int ne02,
190+
ulong nb00,
191+
ulong nb01,
192+
ulong nb02,
193+
ulong nb03,
194+
int ne10,
195+
int ne11,
196+
int ne12,
197+
ulong nb10,
198+
ulong nb11,
199+
ulong nb12,
200+
ulong nb13,
201+
int ne0,
202+
int ne1,
203+
int r2,
204+
int r3
205+
) {
206+
src0 = (global char*)((global char*)src0 + offset0);
207+
src1 = (global char*)((global char*)src1 + offset1);
208+
dst = (global float*)((global char*)dst + offsetd);
209+
210+
const int r0_base = get_group_id(0) * 2;
211+
const int kv_grp = get_group_id(2); // KV head group; im = kv_grp*r2 + q
212+
213+
const int i12_kv = kv_grp % ne02;
214+
const int i13_kv = kv_grp / ne02;
215+
216+
const int lid = get_sub_group_local_id();
217+
const int subhalf = lid >> 5; // 0 or 1 (which K row in the WG)
218+
const int intra = lid & 31; // 0..31 (lane within the half)
219+
220+
const int r0 = r0_base + subhalf;
221+
const int r0c = r0 < ne01 ? r0 : 0; // clamp OOB to row 0; skip write below
222+
223+
// K row pointer for this lane (one K row per half-wave).
224+
const ulong k_off = (ulong)r0c*nb01 + (ulong)i12_kv*nb02 + (ulong)i13_kv*nb03;
225+
global half4 * x4 = (global half4 *)(src0 + k_off);
226+
227+
global float4 * y4[MUL_MAT_F16_F32_L4_DR_LS_R2_MAX];
228+
#pragma unroll
229+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
230+
const int i12_q = i12_kv*r2 + q;
231+
const ulong q_off = (ulong)i12_q*nb12 + (ulong)i13_kv*nb13;
232+
y4[q] = (global float4 *)(src1 + q_off);
233+
}
234+
235+
float partial[MUL_MAT_F16_F32_L4_DR_LS_R2_MAX];
236+
#pragma unroll
237+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
238+
partial[q] = 0.0f;
239+
}
240+
241+
const int n_chunks = ne00 / 4;
242+
243+
for (int i = intra; i < n_chunks; i += 32) {
244+
float4 k = convert_float4(x4[i]);
245+
246+
#pragma unroll
247+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
248+
if (q < r2) {
249+
float4 v = y4[q][i];
250+
partial[q] = mad(k.s0, v.s0, partial[q]);
251+
partial[q] = mad(k.s1, v.s1, partial[q]);
252+
partial[q] = mad(k.s2, v.s2, partial[q]);
253+
partial[q] = mad(k.s3, v.s3, partial[q]);
254+
}
255+
}
256+
}
257+
258+
// half-wave reduction
259+
#pragma unroll
260+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
261+
if (q < r2) {
262+
partial[q] += sub_group_shuffle_xor(partial[q], 1u);
263+
partial[q] += sub_group_shuffle_xor(partial[q], 2u);
264+
partial[q] += sub_group_shuffle_xor(partial[q], 4u);
265+
partial[q] += sub_group_shuffle_xor(partial[q], 8u);
266+
partial[q] += sub_group_shuffle_xor(partial[q], 16u);
267+
}
268+
}
269+
270+
if (intra == 0 && r0 < ne01) {
271+
#pragma unroll
272+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
273+
if (q < r2) {
274+
const int im = i12_kv*r2 + q + i13_kv*ne12;
275+
dst[im*ne1*ne0 + r0] = partial[q];
276+
}
277+
}
278+
}
279+
}
280+
281+
REQD_SUBGROUP_SIZE_64
282+
kernel void kernel_mul_mat_f16_f32_l4_dr_lq(
283+
global char * src0,
284+
ulong offset0,
285+
global char * src1,
286+
ulong offset1,
287+
global float * dst,
288+
ulong offsetd,
289+
int ne00,
290+
int ne01,
291+
int ne02,
292+
ulong nb00,
293+
ulong nb01,
294+
ulong nb02,
295+
ulong nb03,
296+
int ne10,
297+
int ne11,
298+
int ne12,
299+
ulong nb10,
300+
ulong nb11,
301+
ulong nb12,
302+
ulong nb13,
303+
int ne0,
304+
int ne1,
305+
int r2,
306+
int r3
307+
) {
308+
src0 = (global char*)((global char*)src0 + offset0);
309+
src1 = (global char*)((global char*)src1 + offset1);
310+
dst = (global float*)((global char*)dst + offsetd);
311+
312+
const int r0_base = get_group_id(0) * 4;
313+
const int kv_grp = get_group_id(2);
314+
315+
const int i12_kv = kv_grp % ne02;
316+
const int i13_kv = kv_grp / ne02;
317+
318+
const int lid = get_sub_group_local_id();
319+
const int subq = lid >> 4; // 0..3 (which K row)
320+
const int intra = lid & 15; // 0..15 (lane within quarter)
321+
322+
const int r0 = r0_base + subq;
323+
const int r0c = r0 < ne01 ? r0 : 0;
324+
325+
const ulong k_off = (ulong)r0c*nb01 + (ulong)i12_kv*nb02 + (ulong)i13_kv*nb03;
326+
global half4 * x4 = (global half4 *)(src0 + k_off);
327+
328+
global float4 * y4[MUL_MAT_F16_F32_L4_DR_LS_R2_MAX];
329+
#pragma unroll
330+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
331+
const int i12_q = i12_kv*r2 + q;
332+
const ulong q_off = (ulong)i12_q*nb12 + (ulong)i13_kv*nb13;
333+
y4[q] = (global float4 *)(src1 + q_off);
334+
}
335+
336+
float partial[MUL_MAT_F16_F32_L4_DR_LS_R2_MAX];
337+
#pragma unroll
338+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
339+
partial[q] = 0.0f;
340+
}
341+
342+
const int n_chunks = ne00 / 4;
343+
344+
for (int i = intra; i < n_chunks; i += 16) {
345+
float4 k = convert_float4(x4[i]);
346+
347+
#pragma unroll
348+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
349+
if (q < r2) {
350+
float4 v = y4[q][i];
351+
partial[q] = mad(k.s0, v.s0, partial[q]);
352+
partial[q] = mad(k.s1, v.s1, partial[q]);
353+
partial[q] = mad(k.s2, v.s2, partial[q]);
354+
partial[q] = mad(k.s3, v.s3, partial[q]);
355+
}
356+
}
357+
}
358+
359+
// quarter-wave reduction
360+
#pragma unroll
361+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
362+
if (q < r2) {
363+
partial[q] += sub_group_shuffle_xor(partial[q], 1u);
364+
partial[q] += sub_group_shuffle_xor(partial[q], 2u);
365+
partial[q] += sub_group_shuffle_xor(partial[q], 4u);
366+
partial[q] += sub_group_shuffle_xor(partial[q], 8u);
367+
}
368+
}
369+
370+
if (intra == 0 && r0 < ne01) {
371+
#pragma unroll
372+
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
373+
if (q < r2) {
374+
const int im = i12_kv*r2 + q + i13_kv*ne12;
375+
dst[im*ne1*ne0 + r0] = partial[q];
376+
}
377+
}
378+
}
379+
}
380+
#endif // ADRENO_GPU

0 commit comments

Comments
 (0)