Hongqiang Wang
c3d47e696b
* opencl: route `sub_group_shuffle_xor` to qcom ext when KHR ext is unavailable KHR `sub_group_shuffle_xor` is not defined by compiler when `cl_qcom_subgroup_shuffle` is present, causing certain FA kernels fail to build. Define the KHR shuffle_xor using the qcom extension. * opencl: skip FA kernels with mixed and quant types for A7x to avoid compiler crash
1238 lines
37 KiB
Common Lisp
1238 lines
37 KiB
Common Lisp
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
|
||
|
||
#ifdef cl_intel_subgroups
|
||
#pragma OPENCL EXTENSION cl_intel_subgroups : enable
|
||
#else
|
||
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
|
||
#endif
|
||
|
||
#ifdef cl_intel_required_subgroup_size
|
||
#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
|
||
#define INTEL_GPU 1
|
||
#define REQD_SUBGROUP_SIZE_16 __attribute__((intel_reqd_sub_group_size(16)))
|
||
#define REQD_SUBGROUP_SIZE_32 __attribute__((intel_reqd_sub_group_size(32)))
|
||
#elif defined(cl_qcom_reqd_sub_group_size)
|
||
#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
|
||
#define ADRENO_GPU 1
|
||
#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
|
||
#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
|
||
#endif
|
||
|
||
#ifdef cl_khr_subgroup_shuffle
|
||
#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
|
||
#define HAS_SUBGROUP_SHUFFLE 1
|
||
#elif defined(cl_qcom_subgroup_shuffle)
|
||
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
|
||
#define HAS_SUBGROUP_SHUFFLE 1
|
||
// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
|
||
// name, so calling it is an implicit declaration and the program fails to build.
|
||
// Route it to the qcom builtin.
|
||
#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
|
||
#endif
|
||
|
||
// Assumes row size (ne00) is a multiple of 4
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char*)((global char*)src0 + offset0);
|
||
src1 = (global char*)((global char*)src1 + offset1);
|
||
dst = (global float*)((global char*)dst + offsetd);
|
||
|
||
int nrows = ne11;
|
||
int r0 = get_group_id(0);
|
||
int im = get_group_id(2);
|
||
|
||
int i12 = im%ne12;
|
||
int i13 = im/ne12;
|
||
|
||
ulong offset_src0 = r0*nb01 + (i12/r2)*nb02 + (i13/r3)*nb03;
|
||
|
||
global half4 * x4 = (global half4 *) (src0 + offset_src0);
|
||
|
||
for (int r1 = 0; r1 < nrows; ++r1) {
|
||
ulong offset_src1 = r1*nb11 + (i12 )*nb12 + (i13 )*nb13;
|
||
|
||
global float4 * y4 = (global float4 *) (src1 + offset_src1);
|
||
|
||
float sumf = 0;
|
||
for (int i = get_sub_group_local_id(); i < ne00/4; i += get_max_sub_group_size()) {
|
||
sumf += convert_float(x4[i].s0) * y4[i].s0;
|
||
sumf += convert_float(x4[i].s1) * y4[i].s1;
|
||
sumf += convert_float(x4[i].s2) * y4[i].s2;
|
||
sumf += convert_float(x4[i].s3) * y4[i].s3;
|
||
}
|
||
|
||
float all_sum = sub_group_reduce_add(sumf);
|
||
if (get_sub_group_local_id() == 0) {
|
||
dst[im*ne1*ne0 + r1*ne0 + r0] = all_sum;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Each subgroup produces DR_NDST outputs, assumes ne11 == 1
|
||
#define MUL_MAT_F16_F32_L4_DR_NDST 4
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_dr(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char*)((global char*)src0 + offset0);
|
||
src1 = (global char*)((global char*)src1 + offset1);
|
||
dst = (global float*)((global char*)dst + offsetd);
|
||
|
||
const int r0_base = get_group_id(0) * MUL_MAT_F16_F32_L4_DR_NDST;
|
||
const int im = get_group_id(2);
|
||
|
||
const int i12 = im % ne12;
|
||
const int i13 = im / ne12;
|
||
|
||
// assume ne11 == 1
|
||
const ulong offset_src1 = i12*nb12 + i13*nb13;
|
||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||
|
||
global half4 * x4[MUL_MAT_F16_F32_L4_DR_NDST];
|
||
float sumf[MUL_MAT_F16_F32_L4_DR_NDST];
|
||
|
||
const ulong k_head_off = (i12/r2)*nb02 + (i13/r3)*nb03;
|
||
|
||
#pragma unroll
|
||
for (int n = 0; n < MUL_MAT_F16_F32_L4_DR_NDST; ++n) {
|
||
int r0 = r0_base + n;
|
||
int r0c = r0 < ne01 ? r0 : 0;
|
||
ulong off = (ulong)r0c*nb01 + k_head_off;
|
||
x4[n] = (global half4 *)(src0 + off);
|
||
sumf[n] = 0.0f;
|
||
}
|
||
|
||
const int n_chunks = ne00 / 4;
|
||
const int sg_size = get_max_sub_group_size();
|
||
const int lid = get_sub_group_local_id();
|
||
|
||
for (int i = lid; i < n_chunks; i += sg_size) {
|
||
float4 q = y4[i];
|
||
#pragma unroll
|
||
for (int n = 0; n < MUL_MAT_F16_F32_L4_DR_NDST; ++n) {
|
||
float4 k = convert_float4(x4[n][i]);
|
||
sumf[n] = mad(k.s0, q.s0, sumf[n]);
|
||
sumf[n] = mad(k.s1, q.s1, sumf[n]);
|
||
sumf[n] = mad(k.s2, q.s2, sumf[n]);
|
||
sumf[n] = mad(k.s3, q.s3, sumf[n]);
|
||
}
|
||
}
|
||
|
||
#pragma unroll
|
||
for (int n = 0; n < MUL_MAT_F16_F32_L4_DR_NDST; ++n) {
|
||
float reduced = sub_group_reduce_add(sumf[n]);
|
||
int r0 = r0_base + n;
|
||
if (lid == 0 && r0 < ne01) {
|
||
dst[im*ne1*ne0 + r0] = reduced;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Kernels for decoding, Adreno only for now
|
||
#define MUL_MAT_F16_F32_L4_DR_LS_R2_MAX 8
|
||
|
||
#ifdef ADRENO_GPU
|
||
#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
|
||
#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
|
||
|
||
REQD_SUBGROUP_SIZE_64
|
||
kernel void kernel_mul_mat_f16_f32_l4_dr_ls(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char*)((global char*)src0 + offset0);
|
||
src1 = (global char*)((global char*)src1 + offset1);
|
||
dst = (global float*)((global char*)dst + offsetd);
|
||
|
||
const int r0_base = get_group_id(0) * 2;
|
||
const int kv_grp = get_group_id(2); // KV head group; im = kv_grp*r2 + q
|
||
|
||
const int i12_kv = kv_grp % ne02;
|
||
const int i13_kv = kv_grp / ne02;
|
||
|
||
const int lid = get_sub_group_local_id();
|
||
const int subhalf = lid >> 5; // 0 or 1 (which K row in the WG)
|
||
const int intra = lid & 31; // 0..31 (lane within the half)
|
||
|
||
const int r0 = r0_base + subhalf;
|
||
const int r0c = r0 < ne01 ? r0 : 0; // clamp OOB to row 0; skip write below
|
||
|
||
// K row pointer for this lane (one K row per half-wave).
|
||
const ulong k_off = (ulong)r0c*nb01 + (ulong)i12_kv*nb02 + (ulong)i13_kv*nb03;
|
||
global half4 * x4 = (global half4 *)(src0 + k_off);
|
||
|
||
global float4 * y4[MUL_MAT_F16_F32_L4_DR_LS_R2_MAX];
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
const int i12_q = i12_kv*r2 + q;
|
||
const ulong q_off = (ulong)i12_q*nb12 + (ulong)i13_kv*nb13;
|
||
y4[q] = (global float4 *)(src1 + q_off);
|
||
}
|
||
|
||
float partial[MUL_MAT_F16_F32_L4_DR_LS_R2_MAX];
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
partial[q] = 0.0f;
|
||
}
|
||
|
||
const int n_chunks = ne00 / 4;
|
||
|
||
for (int i = intra; i < n_chunks; i += 32) {
|
||
float4 k = convert_float4(x4[i]);
|
||
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
if (q < r2) {
|
||
float4 v = y4[q][i];
|
||
partial[q] = mad(k.s0, v.s0, partial[q]);
|
||
partial[q] = mad(k.s1, v.s1, partial[q]);
|
||
partial[q] = mad(k.s2, v.s2, partial[q]);
|
||
partial[q] = mad(k.s3, v.s3, partial[q]);
|
||
}
|
||
}
|
||
}
|
||
|
||
// half-wave reduction
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
if (q < r2) {
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 1u);
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 2u);
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 4u);
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 8u);
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 16u);
|
||
}
|
||
}
|
||
|
||
if (intra == 0 && r0 < ne01) {
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
if (q < r2) {
|
||
const int im = i12_kv*r2 + q + i13_kv*ne12;
|
||
dst[im*ne1*ne0 + r0] = partial[q];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
REQD_SUBGROUP_SIZE_64
|
||
kernel void kernel_mul_mat_f16_f32_l4_dr_lq(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char*)((global char*)src0 + offset0);
|
||
src1 = (global char*)((global char*)src1 + offset1);
|
||
dst = (global float*)((global char*)dst + offsetd);
|
||
|
||
const int r0_base = get_group_id(0) * 4;
|
||
const int kv_grp = get_group_id(2);
|
||
|
||
const int i12_kv = kv_grp % ne02;
|
||
const int i13_kv = kv_grp / ne02;
|
||
|
||
const int lid = get_sub_group_local_id();
|
||
const int subq = lid >> 4; // 0..3 (which K row)
|
||
const int intra = lid & 15; // 0..15 (lane within quarter)
|
||
|
||
const int r0 = r0_base + subq;
|
||
const int r0c = r0 < ne01 ? r0 : 0;
|
||
|
||
const ulong k_off = (ulong)r0c*nb01 + (ulong)i12_kv*nb02 + (ulong)i13_kv*nb03;
|
||
global half4 * x4 = (global half4 *)(src0 + k_off);
|
||
|
||
global float4 * y4[MUL_MAT_F16_F32_L4_DR_LS_R2_MAX];
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
const int i12_q = i12_kv*r2 + q;
|
||
const ulong q_off = (ulong)i12_q*nb12 + (ulong)i13_kv*nb13;
|
||
y4[q] = (global float4 *)(src1 + q_off);
|
||
}
|
||
|
||
float partial[MUL_MAT_F16_F32_L4_DR_LS_R2_MAX];
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
partial[q] = 0.0f;
|
||
}
|
||
|
||
const int n_chunks = ne00 / 4;
|
||
|
||
for (int i = intra; i < n_chunks; i += 16) {
|
||
float4 k = convert_float4(x4[i]);
|
||
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
if (q < r2) {
|
||
float4 v = y4[q][i];
|
||
partial[q] = mad(k.s0, v.s0, partial[q]);
|
||
partial[q] = mad(k.s1, v.s1, partial[q]);
|
||
partial[q] = mad(k.s2, v.s2, partial[q]);
|
||
partial[q] = mad(k.s3, v.s3, partial[q]);
|
||
}
|
||
}
|
||
}
|
||
|
||
// quarter-wave reduction
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
if (q < r2) {
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 1u);
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 2u);
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 4u);
|
||
partial[q] += sub_group_shuffle_xor(partial[q], 8u);
|
||
}
|
||
}
|
||
|
||
if (intra == 0 && r0 < ne01) {
|
||
#pragma unroll
|
||
for (int q = 0; q < MUL_MAT_F16_F32_L4_DR_LS_R2_MAX; ++q) {
|
||
if (q < r2) {
|
||
const int im = i12_kv*r2 + q + i13_kv*ne12;
|
||
dst[im*ne1*ne0 + r0] = partial[q];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
#endif // ADRENO_GPU
|
||
|
||
#define N_ROWS_PER_WG 8
|
||
#define N_OUTS_PER_WG 8
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_x8(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char *)((global char *)src0 + offset0);
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int sgs_sz = get_max_sub_group_size();
|
||
|
||
const int r0_base = get_group_id(0) * N_ROWS_PER_WG;
|
||
const int im = get_group_id(2);
|
||
|
||
const int i12 = im % ne12;
|
||
const int i13 = im / ne12;
|
||
|
||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||
|
||
__local float4 q_loc[64]; // ne00/4 max for sub_group_size 64
|
||
if (sgs_lid < ne00 / 4) {
|
||
q_loc[sgs_lid] = y4[sgs_lid];
|
||
}
|
||
barrier(CLK_LOCAL_MEM_FENCE);
|
||
|
||
#pragma unroll
|
||
for (int dr = 0; dr < N_ROWS_PER_WG; ++dr) {
|
||
const int r0 = r0_base + dr;
|
||
if (r0 >= ne01) return;
|
||
|
||
const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||
|
||
float sumf = 0.0f;
|
||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||
const half4 k4 = x4[i];
|
||
const float4 q = q_loc[i];
|
||
sumf += convert_float(k4.s0) * q.s0
|
||
+ convert_float(k4.s1) * q.s1
|
||
+ convert_float(k4.s2) * q.s2
|
||
+ convert_float(k4.s3) * q.s3;
|
||
}
|
||
|
||
const float all_sum = sub_group_reduce_add(sumf);
|
||
if (sgs_lid == 0) {
|
||
dst[im * ne1 * ne0 + r0] = all_sum; // ne11 == 1, so r1==0
|
||
}
|
||
}
|
||
}
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_y8(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char *)((global char *)src0 + offset0);
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int sgs_sz = get_max_sub_group_size();
|
||
|
||
const int r0_base = get_group_id(0) * N_OUTS_PER_WG;
|
||
const int im = get_group_id(2);
|
||
|
||
const int i12 = im % ne12;
|
||
const int i13 = im / ne12;
|
||
|
||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||
|
||
global half4 * x4_o[N_OUTS_PER_WG];
|
||
#pragma unroll
|
||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||
const int r0 = r0_base + o;
|
||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||
const ulong off = r0c * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||
x4_o[o] = (global half4 *)(src0 + off);
|
||
}
|
||
|
||
float sum[N_OUTS_PER_WG] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f };
|
||
|
||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||
const float4 q4 = y4[i];
|
||
#pragma unroll
|
||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||
const half4 v4 = x4_o[o][i];
|
||
sum[o] += convert_float(v4.s0) * q4.s0
|
||
+ convert_float(v4.s1) * q4.s1
|
||
+ convert_float(v4.s2) * q4.s2
|
||
+ convert_float(v4.s3) * q4.s3;
|
||
}
|
||
}
|
||
|
||
#pragma unroll
|
||
for (int o = 0; o < N_OUTS_PER_WG; ++o) {
|
||
const int r0 = r0_base + o;
|
||
const float s = sub_group_reduce_add(sum[o]);
|
||
if (sgs_lid == 0 && r0 < ne01) {
|
||
dst[im * ne1 * ne0 + r0] = s;
|
||
}
|
||
}
|
||
}
|
||
|
||
#define N_OUTS_PAIR 8
|
||
#define N_PAIRS_PAIR (N_OUTS_PAIR / 2)
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_x8_pair(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char *)((global char *)src0 + offset0);
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int half_id = sgs_lid >> 5; // 0 = lower half, 1 = upper half
|
||
const int lane_h = sgs_lid & 31; // lane 0..31 within half
|
||
|
||
const int r0_base = get_group_id(0) * N_OUTS_PAIR;
|
||
const int im = get_group_id(2);
|
||
|
||
const int i12 = im % ne12;
|
||
const int i13 = im / ne12;
|
||
|
||
const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
|
||
global float4 * y4 = (global float4 *)(src1 + offset_src1);
|
||
|
||
__local float4 q_loc[64]; // ne00/4 max for sub_group_size 64
|
||
if (sgs_lid < ne00 / 4) {
|
||
q_loc[sgs_lid] = y4[sgs_lid];
|
||
}
|
||
barrier(CLK_LOCAL_MEM_FENCE);
|
||
|
||
const int dk_vec = ne00 / 4;
|
||
|
||
#pragma unroll
|
||
for (int p = 0; p < N_PAIRS_PAIR; ++p) {
|
||
const int r0 = r0_base + 2 * p + half_id;
|
||
|
||
const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
|
||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||
|
||
float sumf = 0.0f;
|
||
for (int i = lane_h; i < dk_vec; i += 32) {
|
||
const half4 k4 = x4[i];
|
||
const float4 q = q_loc[i];
|
||
sumf += convert_float(k4.s0) * q.s0
|
||
+ convert_float(k4.s1) * q.s1
|
||
+ convert_float(k4.s2) * q.s2
|
||
+ convert_float(k4.s3) * q.s3;
|
||
}
|
||
|
||
sumf += sub_group_shuffle_xor(sumf, 16);
|
||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||
|
||
if (lane_h == 0) {
|
||
dst[im * ne1 * ne0 + r0] = sumf;
|
||
}
|
||
}
|
||
}
|
||
|
||
#define N_K_ROWS_GQA 16
|
||
#define GQA_RATIO_GQA 8
|
||
#define LANES_PER_QH 8 // 64 / GQA_RATIO_GQA
|
||
#define DK_VEC_GQA 32 // DK / 4 for DK=128
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char *)((global char *)src0 + offset0);
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int q_id = sgs_lid >> 3; // 0..7: which Q-head (8 per WG)
|
||
const int lane_q = sgs_lid & 7; // 0..7: lane within Q-head partition
|
||
|
||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
|
||
const int im_kv = get_group_id(2);
|
||
|
||
const int i02 = im_kv % ne02; // K-head index (also K2 batch)
|
||
const int i03 = im_kv / ne02; // n13 batch index
|
||
|
||
const int q_head_lo = i02 * GQA_RATIO_GQA;
|
||
|
||
__local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA]; // 4 × 32 = 128 float4
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
|
||
const int qh_idx = q_head_lo + qh;
|
||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||
|
||
if (sgs_lid < DK_VEC_GQA) {
|
||
q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
|
||
}
|
||
}
|
||
barrier(CLK_LOCAL_MEM_FENCE);
|
||
|
||
// K base offset for this WG. All 8 K-rows × 4 Q-heads share this K-head.
|
||
const ulong offset_src0_base = (i02) * nb02 + (i03 / r3) * nb03;
|
||
|
||
#pragma unroll
|
||
for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
|
||
const int r0 = r0_base + dr;
|
||
|
||
const ulong offset_src0 = r0 * nb01 + offset_src0_base;
|
||
global half4 * x4 = (global half4 *)(src0 + offset_src0);
|
||
|
||
float sumf = 0.0f;
|
||
#pragma unroll
|
||
for (int t = 0; t < 4; ++t) {
|
||
const int i = lane_q + t * LANES_PER_QH; // 8, 16, 24-step
|
||
const half4 k4 = x4[i];
|
||
const float4 q = q_loc[q_id * DK_VEC_GQA + i];
|
||
sumf += convert_float(k4.s0) * q.s0
|
||
+ convert_float(k4.s1) * q.s1
|
||
+ convert_float(k4.s2) * q.s2
|
||
+ convert_float(k4.s3) * q.s3;
|
||
}
|
||
|
||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||
|
||
if (lane_q == 0) {
|
||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||
}
|
||
}
|
||
}
|
||
|
||
#define N_DV_ROWS_Y8GQA 8
|
||
#define GQA_RATIO_Y8GQA 8
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_y8_gqa(
|
||
global char * src0,
|
||
ulong offset0,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb00,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src0 = (global char *)((global char *)src0 + offset0);
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int sgs_sz = get_max_sub_group_size();
|
||
|
||
const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
|
||
const int im_kv = get_group_id(2);
|
||
|
||
const int i02 = im_kv % ne02; // K-head index
|
||
const int i03 = im_kv / ne02; // n13 batch index
|
||
|
||
// GQA Q-heads sharing this K-head.
|
||
const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
|
||
|
||
global float4 * y4_q[GQA_RATIO_Y8GQA];
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||
const int qh_idx = q_head_lo + qh;
|
||
y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||
}
|
||
|
||
global half4 * x4_o[N_DV_ROWS_Y8GQA];
|
||
#pragma unroll
|
||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||
const int r0 = r0_base + o;
|
||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||
const ulong off = r0c * nb01 + (i02) * nb02 + (i03 / r3) * nb03;
|
||
x4_o[o] = (global half4 *)(src0 + off);
|
||
}
|
||
|
||
float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
|
||
|
||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||
// load 8 V values (one per DV row), same K-head, K-pos = i.
|
||
half4 v[N_DV_ROWS_Y8GQA];
|
||
#pragma unroll
|
||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||
v[o] = x4_o[o][i];
|
||
}
|
||
|
||
// load 8 softmax values (one per Q-head).
|
||
float4 q[GQA_RATIO_Y8GQA];
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||
q[qh] = y4_q[qh][i];
|
||
}
|
||
|
||
#pragma unroll
|
||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||
const float4 vf = (float4)(convert_float(v[o].s0),
|
||
convert_float(v[o].s1),
|
||
convert_float(v[o].s2),
|
||
convert_float(v[o].s3));
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||
sum[o][qh] += vf.s0 * q[qh].s0
|
||
+ vf.s1 * q[qh].s1
|
||
+ vf.s2 * q[qh].s2
|
||
+ vf.s3 * q[qh].s3;
|
||
}
|
||
}
|
||
}
|
||
|
||
#pragma unroll
|
||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||
const int r0 = r0_base + o;
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||
const float s = sub_group_reduce_add(sum[o][qh]);
|
||
if (sgs_lid == 0 && r0 < ne01) {
|
||
const int im_out = i03 * ne12 + (q_head_lo + qh);
|
||
dst[im_out * ne1 * ne0 + r0] = s;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4_img(
|
||
__read_only image1d_buffer_t src0_img,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int q_id = sgs_lid >> 3; // 0..7: which Q-head (8 per WG)
|
||
const int lane_q = sgs_lid & 7; // 0..7: lane within Q-head partition
|
||
|
||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
|
||
const int im_kv = get_group_id(2);
|
||
|
||
const int i02 = im_kv % ne02;
|
||
const int i03 = im_kv / ne02;
|
||
|
||
const int q_head_lo = i02 * GQA_RATIO_GQA;
|
||
|
||
__local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA];
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
|
||
const int qh_idx = q_head_lo + qh;
|
||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||
if (sgs_lid < DK_VEC_GQA) {
|
||
q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
|
||
}
|
||
}
|
||
barrier(CLK_LOCAL_MEM_FENCE);
|
||
|
||
const int pitch_px_row = (int)(nb01 >> 4);
|
||
const int pitch_px_head = (int)(nb02 >> 4);
|
||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||
|
||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||
|
||
#pragma unroll
|
||
for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
|
||
const int r0 = r0_base + dr;
|
||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||
|
||
float sumf = 0.0f;
|
||
#pragma unroll
|
||
for (int t = 0; t < 2; ++t) {
|
||
const int p = lane_q + t * LANES_PER_QH; // pixel idx in row, 0..15
|
||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||
const int i0 = 2 * p; // first half4 idx
|
||
const float4 qa = q_loc[q_id * DK_VEC_GQA + i0 ];
|
||
const float4 qb = q_loc[q_id * DK_VEC_GQA + i0 + 1];
|
||
sumf += convert_float(k8.s0) * qa.s0
|
||
+ convert_float(k8.s1) * qa.s1
|
||
+ convert_float(k8.s2) * qa.s2
|
||
+ convert_float(k8.s3) * qa.s3
|
||
+ convert_float(k8.s4) * qb.s0
|
||
+ convert_float(k8.s5) * qb.s1
|
||
+ convert_float(k8.s6) * qb.s2
|
||
+ convert_float(k8.s7) * qb.s3;
|
||
}
|
||
|
||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||
|
||
if (lane_q == 0) {
|
||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||
}
|
||
}
|
||
}
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_y8_gqa_img(
|
||
__read_only image1d_buffer_t src0_img,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int sgs_sz = get_max_sub_group_size();
|
||
|
||
const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
|
||
const int im_kv = get_group_id(2);
|
||
|
||
const int i02 = im_kv % ne02;
|
||
const int i03 = im_kv / ne02;
|
||
|
||
const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
|
||
|
||
// Q (= softmax(KQ)) base pointers per Q-head
|
||
global float4 * y4_q[GQA_RATIO_Y8GQA];
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||
const int qh_idx = q_head_lo + qh;
|
||
y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||
}
|
||
|
||
const int pitch_px_row = (int)(nb01 >> 3);
|
||
const int pitch_px_head = (int)(nb02 >> 3);
|
||
const int pitch_px_n13 = (int)(nb03 >> 3);
|
||
|
||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||
|
||
// per-DV-row pixel base
|
||
int row_px_base[N_DV_ROWS_Y8GQA];
|
||
#pragma unroll
|
||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||
const int r0 = r0_base + o;
|
||
const int r0c = (r0 < ne01) ? r0 : 0;
|
||
row_px_base[o] = r0c * pitch_px_row + head_px_base;
|
||
}
|
||
|
||
float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
|
||
|
||
for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
|
||
half4 v[N_DV_ROWS_Y8GQA];
|
||
|
||
#pragma unroll
|
||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||
v[o] = read_imageh(src0_img, row_px_base[o] + i);
|
||
}
|
||
|
||
float4 q[GQA_RATIO_Y8GQA];
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||
q[qh] = y4_q[qh][i];
|
||
}
|
||
// 64 mads.
|
||
#pragma unroll
|
||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||
const float4 vf = (float4)(convert_float(v[o].s0),
|
||
convert_float(v[o].s1),
|
||
convert_float(v[o].s2),
|
||
convert_float(v[o].s3));
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||
sum[o][qh] += vf.s0 * q[qh].s0
|
||
+ vf.s1 * q[qh].s1
|
||
+ vf.s2 * q[qh].s2
|
||
+ vf.s3 * q[qh].s3;
|
||
}
|
||
}
|
||
}
|
||
|
||
#pragma unroll
|
||
for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
|
||
const int r0 = r0_base + o;
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
|
||
const float s = sub_group_reduce_add(sum[o][qh]);
|
||
if (sgs_lid == 0 && r0 < ne01) {
|
||
const int im_out = i03 * ne12 + (q_head_lo + qh);
|
||
dst[im_out * ne1 * ne0 + r0] = s;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#define N_K_ROWS_GQA_R4 16
|
||
#define GQA_RATIO_R4 4
|
||
#define LANES_PER_QH_R4 16 // = 64 / GQA_RATIO_R4
|
||
#define DK_VEC_R4 32 // DK / 4 for DK=128
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img(
|
||
__read_only image1d_buffer_t src0_img,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int q_id = sgs_lid >> 4; // 0..3
|
||
const int lane_q = sgs_lid & 15; // 0..15
|
||
|
||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R4;
|
||
const int im_kv = get_group_id(2);
|
||
|
||
const int i02 = im_kv % ne02;
|
||
const int i03 = im_kv / ne02;
|
||
|
||
const int q_head_lo = i02 * GQA_RATIO_R4;
|
||
|
||
__local float4 q_loc[GQA_RATIO_R4 * DK_VEC_R4];
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_R4; ++qh) {
|
||
const int qh_idx = q_head_lo + qh;
|
||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||
if (sgs_lid < DK_VEC_R4) {
|
||
q_loc[qh * DK_VEC_R4 + sgs_lid] = y4[sgs_lid];
|
||
}
|
||
}
|
||
barrier(CLK_LOCAL_MEM_FENCE);
|
||
|
||
const int pitch_px_row = (int)(nb01 >> 4);
|
||
const int pitch_px_head = (int)(nb02 >> 4);
|
||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||
|
||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||
|
||
#pragma unroll
|
||
for (int dr = 0; dr < N_K_ROWS_GQA_R4; ++dr) {
|
||
const int r0 = r0_base + dr;
|
||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||
|
||
const int p = lane_q;
|
||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||
const int i0 = 2 * p;
|
||
const float4 qa = q_loc[q_id * DK_VEC_R4 + i0 ];
|
||
const float4 qb = q_loc[q_id * DK_VEC_R4 + i0 + 1];
|
||
|
||
float sumf =
|
||
convert_float(k8.s0) * qa.s0
|
||
+ convert_float(k8.s1) * qa.s1
|
||
+ convert_float(k8.s2) * qa.s2
|
||
+ convert_float(k8.s3) * qa.s3
|
||
+ convert_float(k8.s4) * qb.s0
|
||
+ convert_float(k8.s5) * qb.s1
|
||
+ convert_float(k8.s6) * qb.s2
|
||
+ convert_float(k8.s7) * qb.s3;
|
||
|
||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||
|
||
if (lane_q == 0) {
|
||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||
}
|
||
}
|
||
}
|
||
|
||
#define N_K_ROWS_GQA_R2_DK256 16
|
||
#define GQA_RATIO_R2 2
|
||
#define LANES_PER_QH_R2 32 // = 64 / GQA_RATIO_R2
|
||
#define DK_VEC_DK256 64 // DK / 4 for DK=256
|
||
|
||
#ifdef ADRENO_GPU
|
||
REQD_SUBGROUP_SIZE_64
|
||
#endif
|
||
kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img(
|
||
__read_only image1d_buffer_t src0_img,
|
||
global char * src1,
|
||
ulong offset1,
|
||
global float * dst,
|
||
ulong offsetd,
|
||
int ne00,
|
||
int ne01,
|
||
int ne02,
|
||
ulong nb01,
|
||
ulong nb02,
|
||
ulong nb03,
|
||
int ne10,
|
||
int ne11,
|
||
int ne12,
|
||
ulong nb10,
|
||
ulong nb11,
|
||
ulong nb12,
|
||
ulong nb13,
|
||
int ne0,
|
||
int ne1,
|
||
int r2,
|
||
int r3
|
||
) {
|
||
src1 = (global char *)((global char *)src1 + offset1);
|
||
dst = (global float*)((global char *)dst + offsetd);
|
||
|
||
const int sgs_lid = get_sub_group_local_id();
|
||
const int q_id = sgs_lid >> 5; // 0..1
|
||
const int lane_q = sgs_lid & 31; // 0..31
|
||
|
||
const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R2_DK256;
|
||
const int im_kv = get_group_id(2);
|
||
|
||
const int i02 = im_kv % ne02;
|
||
const int i03 = im_kv / ne02;
|
||
|
||
const int q_head_lo = i02 * GQA_RATIO_R2;
|
||
|
||
__local float4 q_loc[GQA_RATIO_R2 * DK_VEC_DK256];
|
||
#pragma unroll
|
||
for (int qh = 0; qh < GQA_RATIO_R2; ++qh) {
|
||
const int qh_idx = q_head_lo + qh;
|
||
global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
|
||
q_loc[qh * DK_VEC_DK256 + sgs_lid] = y4[sgs_lid];
|
||
}
|
||
barrier(CLK_LOCAL_MEM_FENCE);
|
||
|
||
const int pitch_px_row = (int)(nb01 >> 4);
|
||
const int pitch_px_head = (int)(nb02 >> 4);
|
||
const int pitch_px_n13 = (int)(nb03 >> 4);
|
||
|
||
const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
|
||
|
||
#pragma unroll
|
||
for (int dr = 0; dr < N_K_ROWS_GQA_R2_DK256; ++dr) {
|
||
const int r0 = r0_base + dr;
|
||
const int row_px_base = r0 * pitch_px_row + head_px_base;
|
||
|
||
const int p = lane_q;
|
||
const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
|
||
const int i0 = 2 * p;
|
||
const float4 qa = q_loc[q_id * DK_VEC_DK256 + i0 ];
|
||
const float4 qb = q_loc[q_id * DK_VEC_DK256 + i0 + 1];
|
||
|
||
float sumf =
|
||
convert_float(k8.s0) * qa.s0
|
||
+ convert_float(k8.s1) * qa.s1
|
||
+ convert_float(k8.s2) * qa.s2
|
||
+ convert_float(k8.s3) * qa.s3
|
||
+ convert_float(k8.s4) * qb.s0
|
||
+ convert_float(k8.s5) * qb.s1
|
||
+ convert_float(k8.s6) * qb.s2
|
||
+ convert_float(k8.s7) * qb.s3;
|
||
|
||
sumf += sub_group_shuffle_xor(sumf, 16);
|
||
sumf += sub_group_shuffle_xor(sumf, 8);
|
||
sumf += sub_group_shuffle_xor(sumf, 4);
|
||
sumf += sub_group_shuffle_xor(sumf, 2);
|
||
sumf += sub_group_shuffle_xor(sumf, 1);
|
||
|
||
if (lane_q == 0) {
|
||
const int im_out = i03 * ne12 + (q_head_lo + q_id);
|
||
dst[im_out * ne1 * ne0 + r0] = sumf;
|
||
}
|
||
}
|
||
}
|