#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#pragma OPENCL EXTENSION cl_qcom_subgroup_uniform_load : enable
#pragma OPENCL EXTENSION cl_qcom_subgroup_constant_load : enable
#define bool2 uchar2
#define bool3 uchar3
#define bool4 uchar4
__constant sampler_t smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST__constant sampler_t smp_zero = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST
__kernel void adreno_xmem_attn_q_f32_to_img_scaled(const global void * src_void,
ulong src_offset,
write_only image2d_t dst_image2d,
const float scale,
const int d_head,
const int n_q,
const int n_head,
const int n_head_kv,
const int n_batch,
const ulong src_nb1,
const ulong src_nb2,
const ulong src_nb3) {
const int x = get_global_id(0) const int flat_h = get_global_id(1) const int d = get_global_id(2)
const int heads_total = n_head * n_batch const int kpack = d_head / 4
if (x >= n_q || flat_h >= heads_total || d >= kpack) {
return }
const int batch = flat_h / n_head const int head = flat_h % n_head const int gqa = n_head / n_head_kv const int head_kv = head / gqa const int head_group = head - head_kv * gqa const int compact_h = batch * n_head_kv + head_kv const int compact_x = head_group * n_q + x const int c = d * 4
const global char * src_base = (const global char *) src_void + src_offset const global float * row_ptr = (const global float *) (src_base + batch * src_nb3 + head * src_nb2 + x * src_nb1)
half4 out = (half4) (0.0h) out.x = convert_half(row_ptr[c + 0] * scale) if (c + 1 < d_head) {
out.y = convert_half(row_ptr[c + 1] * scale) }
if (c + 2 < d_head) {
out.z = convert_half(row_ptr[c + 2] * scale) }
if (c + 3 < d_head) {
out.w = convert_half(row_ptr[c + 3] * scale) }
write_imageh(dst_image2d, (int2) (compact_x, compact_h * kpack + d), out)}
__kernel void adreno_xmem_attn_kv_f32_to_img_gqa(const global void * src_void,
ulong src_offset,
write_only image2d_t dst_image2d,
const int d_head,
const int n_kv,
const int n_kv_padded,
const int n_head_kv,
const int n_batch,
const ulong src_nb1,
const ulong src_nb2,
const ulong src_nb3) {
const int x = get_global_id(0) const int flat_h = get_global_id(1) const int d = get_global_id(2)
const int kv_heads_total = n_head_kv * n_batch const int kpack = d_head / 4
if (x >= n_kv_padded || flat_h >= kv_heads_total || d >= kpack) {
return }
const int batch = flat_h / n_head_kv const int head_kv = flat_h % n_head_kv const int c = d * 4
half4 out = (half4) (0.0h) if (x < n_kv) {
const global char * src_base = (const global char *) src_void + src_offset const global float * row_ptr =
(const global float *) (src_base + batch * src_nb3 + head_kv * src_nb2 + x * src_nb1) out.x = convert_half(row_ptr[c + 0]) if (c + 1 < d_head) {
out.y = convert_half(row_ptr[c + 1]) }
if (c + 2 < d_head) {
out.z = convert_half(row_ptr[c + 2]) }
if (c + 3 < d_head) {
out.w = convert_half(row_ptr[c + 3]) }
}
write_imageh(dst_image2d, (int2) (x, flat_h * kpack + d), out)}
__kernel void adreno_xmem_attn_kv_f16_to_img_gqa(const global void * src_void,
ulong src_offset,
write_only image2d_t dst_image2d,
const int d_head,
const int n_kv,
const int n_kv_padded,
const int n_head_kv,
const int n_batch,
const ulong src_nb1,
const ulong src_nb2,
const ulong src_nb3) {
const int x = get_global_id(0) const int flat_h = get_global_id(1) const int d = get_global_id(2)
const int kv_heads_total = n_head_kv * n_batch const int kpack = d_head / 4
if (x >= n_kv_padded || flat_h >= kv_heads_total || d >= kpack) {
return }
const int batch = flat_h / n_head_kv const int head_kv = flat_h % n_head_kv const int c = d * 4
half4 out = (half4) (0.0h) if (x < n_kv) {
const global char * src_base = (const global char *) src_void + src_offset const global half * row_ptr =
(const global half *) (src_base + batch * src_nb3 + head_kv * src_nb2 + x * src_nb1) out.x = row_ptr[c + 0] if (c + 1 < d_head) {
out.y = row_ptr[c + 1] }
if (c + 2 < d_head) {
out.z = row_ptr[c + 2] }
if (c + 3 < d_head) {
out.w = row_ptr[c + 3] }
}
write_imageh(dst_image2d, (int2) (x, flat_h * kpack + d), out)}
__kernel void adreno_xmem_attn_img_to_f32(global void * dst_void,
ulong dst_offset,
read_only image2d_t src_image2d,
const int d_head,
const int n_q,
const int n_head,
const int n_head_kv,
const int n_batch,
const ulong dst_nb1,
const ulong dst_nb2,
const ulong dst_nb3) {
const int x = get_global_id(0) const int flat_h = get_global_id(1) const int d = get_global_id(2)
const int heads_total = n_head * n_batch const int kpack = d_head / 4
if (x >= n_q || flat_h >= heads_total || d >= kpack) {
return }
const int batch = flat_h / n_head const int head = flat_h % n_head const int gqa = n_head / n_head_kv const int head_kv = head / gqa const int head_group = head - head_kv * gqa const int compact_h = batch * n_head_kv + head_kv const int compact_x = head_group * n_q + x const int c = d * 4
global char * dst_base = (global char *) dst_void + dst_offset global float * row_ptr = (global float *) (dst_base + batch * dst_nb3 + x * dst_nb2 + head * dst_nb1)
const half4 in_value = read_imageh(src_image2d, smp_zero, (int2) (compact_x, compact_h * kpack + d)) row_ptr[c + 0] = convert_float(in_value.x) if (c + 1 < d_head) {
row_ptr[c + 1] = convert_float(in_value.y) }
if (c + 2 < d_head) {
row_ptr[c + 2] = convert_float(in_value.z) }
if (c + 3 < d_head) {
row_ptr[c + 3] = convert_float(in_value.w) }
}
__kernel void adreno_xmem_attn_k_gather(global half4 * dst_tensor_buffer,
read_only image2d_t src_tensor_image2d,
const int4 shared_int4_0,
const int4 shared_int4_1) {
int X = get_global_id(0) int Y = get_global_id(1) int S = get_global_id(2) if (X >= shared_int4_0.w || Y >= shared_int4_0.y || S >= shared_int4_0.z) {
return }
half temps[4] temps[0] = (half) (0.f) temps[1] = (half) (0.f) temps[2] = (half) (0.f) temps[3] = (half) (0.f) for (int i = 0 int dst_channel = S * 4 + i if (dst_channel < shared_int4_0.x) {
int s_y = Y int s_x = dst_channel int s_c = X {
int slice_coord_TMP = (s_c) / 4 int sub_ch_coord_TMP = (s_c) % 4 half4 src_TMP = read_imageh(src_tensor_image2d, smp_zero,
(int2) ((s_x), ((s_y) *shared_int4_1.x + (slice_coord_TMP)))) temps[i] = (half[4]){ src_TMP.x, src_TMP.y, src_TMP.z, src_TMP.w }[sub_ch_coord_TMP] } }
}
half4 result result.x = temps[0] result.y = temps[1] result.z = temps[2] result.w = temps[3] dst_tensor_buffer[(((S) *shared_int4_0.y + (Y)) * shared_int4_0.w + (X))] = result}
__kernel void adreno_xmem_attn_pack_k(global half4 * dst_tensor_buffer,
read_only image1d_buffer_t src_image_buffer,
const int4 shared_int4_0,
const int4 shared_int4_1,
const int4 shared_int4_2) {
int linear_index = get_global_id(0) if (linear_index >= shared_int4_0.y) {
return }
if (get_global_id(1) != 0) {
return }
if (get_global_id(2) != 0) {
return }
int dst_o_sp_i_ogroup = linear_index int dst_ogroup = dst_o_sp_i_ogroup % shared_int4_0.x int dst_o_sp_i = dst_o_sp_i_ogroup / shared_int4_0.x int dst_i = dst_o_sp_i % shared_int4_0.z int dst_o_sp = dst_o_sp_i / shared_int4_0.z int dst_sp = dst_o_sp % shared_int4_1.x int dst_o = dst_o_sp / shared_int4_1.x int i_slice = dst_i int o_slice = dst_o * shared_int4_0.x + dst_ogroup int spatial_linear = dst_sp int W = spatial_linear % shared_int4_1.y int H = spatial_linear / shared_int4_1.y half4 w0 = (half4) (0) half4 w1 = (half4) (0) half4 w2 = (half4) (0) half4 w3 = (half4) (0)
if (i_slice * 4 < shared_int4_0.w && o_slice < shared_int4_1.w) {
w0 = read_imageh(src_image_buffer, (((o_slice) *shared_int4_1.z + (W)) * shared_int4_2.x + (i_slice * 4))) }
if (i_slice * 4 + 1 < shared_int4_0.w && o_slice < shared_int4_1.w) {
w1 = read_imageh(src_image_buffer, (((o_slice) *shared_int4_1.z + (W)) * shared_int4_2.x + (i_slice * 4 + 1))) }
if (i_slice * 4 + 2 < shared_int4_0.w && o_slice < shared_int4_1.w) {
w2 = read_imageh(src_image_buffer, (((o_slice) *shared_int4_1.z + (W)) * shared_int4_2.x + (i_slice * 4 + 2))) }
if (i_slice * 4 + 3 < shared_int4_0.w && o_slice < shared_int4_1.w) {
w3 = read_imageh(src_image_buffer, (((o_slice) *shared_int4_1.z + (W)) * shared_int4_2.x + (i_slice * 4 + 3))) }
half4 r0 = w0 half4 r1 = w1 half4 r2 = w2 half4 r3 = w3 dst_tensor_buffer[linear_index * 4 + 0] = r0 dst_tensor_buffer[linear_index * 4 + 1] = r1 dst_tensor_buffer[linear_index * 4 + 2] = r2 dst_tensor_buffer[linear_index * 4 + 3] = r3}
__attribute__((qcom_max_concurrent_subgroups(12))) __kernel void adreno_xmem_attn_qk_gemm(
global half4 * dst_tensor_buffer,
constant half8 * weights_buffer __attribute__((sub_group_uniform)),
constant half8 * xmem_buffer __attribute__((max_constant_size((6144)))),
read_only image2d_t src_tensor_image2d,
const int4 shared_int4_0,
const int4 shared_int4_1,
const int4 shared_int4_2) {
int X = get_group_id(1) * get_local_size(0) + get_local_id(0) int Y = get_group_id(2) * get_local_size(1) + get_local_id(1) int Z = get_group_id(0) * get_local_size(2) + get_local_id(2) if (X >= shared_int4_0.z || Y >= shared_int4_0.x) {
return }
if (Z * 8 >= shared_int4_0.y) {
return }
half4 r0 = (half4) (0.f) half4 r1 = (half4) (0.f) half4 r2 = (half4) (0.f) half4 r3 = (half4) (0.f) half4 r4 = (half4) (0.f) half4 r5 = (half4) (0.f) half4 r6 = (half4) (0.f) half4 r7 = (half4) (0.f) int x_coord = mad24(X, shared_int4_2.y, shared_int4_1.y) int y_coord = mad24(Y, shared_int4_2.z, shared_int4_1.z) int coord_x, coord_y, coord_s int f_offset = (Z * shared_int4_1.w + Y) * shared_int4_1.x * 32
int subgroup_id = (int) ((0x1F & qcom_get_physical_sub_group_id())) subgroup_id = subgroup_id % 12 int c_offset = mul24(subgroup_id, shared_int4_0.w) __constant half16 * weights_cache = (__constant half16 *) &xmem_buffer[c_offset] coord_y = Y coord_x = X coord_s = 0 do {
half4 src0 =
read_imageh(src_tensor_image2d, smp_zero, (int2) ((coord_x), ((coord_y) *shared_int4_2.x + (coord_s)))) coord_s++ half4 src1 =
read_imageh(src_tensor_image2d, smp_zero, (int2) ((coord_x), ((coord_y) *shared_int4_2.x + (coord_s)))) coord_s++ qcom_sub_group_constant_load8(xmem_buffer, weights_buffer, c_offset, f_offset >> 1, 32) f_offset += 64 qcom_sub_group_sync(QCOM_CLK_CONST_LOAD_SYNC) r0 += src0.x * weights_cache[0].s0123 r0 += src0.y * weights_cache[0].s4567 r0 += src0.z * weights_cache[0].s89ab r0 += src0.w * weights_cache[0].scdef r1 += src0.x * weights_cache[1].s0123 r1 += src0.y * weights_cache[1].s4567 r1 += src0.z * weights_cache[1].s89ab r1 += src0.w * weights_cache[1].scdef r2 += src0.x * weights_cache[2].s0123 r2 += src0.y * weights_cache[2].s4567 r2 += src0.z * weights_cache[2].s89ab r2 += src0.w * weights_cache[2].scdef r3 += src0.x * weights_cache[3].s0123 r3 += src0.y * weights_cache[3].s4567 r3 += src0.z * weights_cache[3].s89ab r3 += src0.w * weights_cache[3].scdef r4 += src0.x * weights_cache[4].s0123 r4 += src0.y * weights_cache[4].s4567 r4 += src0.z * weights_cache[4].s89ab r4 += src0.w * weights_cache[4].scdef r5 += src0.x * weights_cache[5].s0123 r5 += src0.y * weights_cache[5].s4567 r5 += src0.z * weights_cache[5].s89ab r5 += src0.w * weights_cache[5].scdef r6 += src0.x * weights_cache[6].s0123 r6 += src0.y * weights_cache[6].s4567 r6 += src0.z * weights_cache[6].s89ab r6 += src0.w * weights_cache[6].scdef r7 += src0.x * weights_cache[7].s0123 r7 += src0.y * weights_cache[7].s4567 r7 += src0.z * weights_cache[7].s89ab r7 += src0.w * weights_cache[7].scdef r0 += src1.x * weights_cache[8].s0123 r0 += src1.y * weights_cache[8].s4567 r0 += src1.z * weights_cache[8].s89ab r0 += src1.w * weights_cache[8].scdef r1 += src1.x * weights_cache[9].s0123 r1 += src1.y * weights_cache[9].s4567 r1 += src1.z * weights_cache[9].s89ab r1 += src1.w * weights_cache[9].scdef r2 += src1.x * weights_cache[10].s0123 r2 += src1.y * weights_cache[10].s4567 r2 += src1.z * weights_cache[10].s89ab r2 += src1.w * weights_cache[10].scdef r3 += src1.x * weights_cache[11].s0123 r3 += src1.y * weights_cache[11].s4567 r3 += src1.z * weights_cache[11].s89ab r3 += src1.w * weights_cache[11].scdef r4 += src1.x * weights_cache[12].s0123 r4 += src1.y * weights_cache[12].s4567 r4 += src1.z * weights_cache[12].s89ab r4 += src1.w * weights_cache[12].scdef r5 += src1.x * weights_cache[13].s0123 r5 += src1.y * weights_cache[13].s4567 r5 += src1.z * weights_cache[13].s89ab r5 += src1.w * weights_cache[13].scdef r6 += src1.x * weights_cache[14].s0123 r6 += src1.y * weights_cache[14].s4567 r6 += src1.z * weights_cache[14].s89ab r6 += src1.w * weights_cache[14].scdef r7 += src1.x * weights_cache[15].s0123 r7 += src1.y * weights_cache[15].s4567 r7 += src1.z * weights_cache[15].s89ab r7 += src1.w * weights_cache[15].scdef } while (coord_s < shared_int4_2.x)
coord_s = mul24(Z, 8) coord_x = X coord_y = Y if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r0) if (coord_s < 0) {
res += read_imageh(src_tensor_image2d, smp_zero, (int2) ((0), ((0) * shared_int4_2.x + (0)))) }
dst_tensor_buffer[(((coord_s) *shared_int4_0.x + (coord_y)) * shared_int4_0.z + (coord_x))] = res coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r1) if (coord_s < 0) {
res += read_imageh(src_tensor_image2d, smp_zero, (int2) ((0), ((0) * shared_int4_2.x + (0)))) }
dst_tensor_buffer[(((coord_s) *shared_int4_0.x + (coord_y)) * shared_int4_0.z + (coord_x))] = res coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r2) if (coord_s < 0) {
res += read_imageh(src_tensor_image2d, smp_zero, (int2) ((0), ((0) * shared_int4_2.x + (0)))) }
dst_tensor_buffer[(((coord_s) *shared_int4_0.x + (coord_y)) * shared_int4_0.z + (coord_x))] = res coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r3) if (coord_s < 0) {
res += read_imageh(src_tensor_image2d, smp_zero, (int2) ((0), ((0) * shared_int4_2.x + (0)))) }
dst_tensor_buffer[(((coord_s) *shared_int4_0.x + (coord_y)) * shared_int4_0.z + (coord_x))] = res coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r4) if (coord_s < 0) {
res += read_imageh(src_tensor_image2d, smp_zero, (int2) ((0), ((0) * shared_int4_2.x + (0)))) }
dst_tensor_buffer[(((coord_s) *shared_int4_0.x + (coord_y)) * shared_int4_0.z + (coord_x))] = res coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r5) if (coord_s < 0) {
res += read_imageh(src_tensor_image2d, smp_zero, (int2) ((0), ((0) * shared_int4_2.x + (0)))) }
dst_tensor_buffer[(((coord_s) *shared_int4_0.x + (coord_y)) * shared_int4_0.z + (coord_x))] = res coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r6) if (coord_s < 0) {
res += read_imageh(src_tensor_image2d, smp_zero, (int2) ((0), ((0) * shared_int4_2.x + (0)))) }
dst_tensor_buffer[(((coord_s) *shared_int4_0.x + (coord_y)) * shared_int4_0.z + (coord_x))] = res coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r7) if (coord_s < 0) {
res += read_imageh(src_tensor_image2d, smp_zero, (int2) ((0), ((0) * shared_int4_2.x + (0)))) }
dst_tensor_buffer[(((coord_s) *shared_int4_0.x + (coord_y)) * shared_int4_0.z + (coord_x))] = res coord_s++ }
}
__kernel void adreno_xmem_attn_softmax_reduce_basic(read_only image1d_buffer_t src_tensor_image_buffer,
write_only image2d_t dst_tensor_image2d,
const int4 shared_int4_0,
const int4 shared_int4_1) {
int X = get_global_id(0) int Y = get_global_id(1) if (X >= shared_int4_0.z || Y >= shared_int4_0.x) {
return }
float sum = 0.0f int end_channel = shared_int4_0.w int end_slice = (end_channel + 3) / 4 int start_channel = 0 int start_slice = start_channel / 4 bool need_per_channels_check = start_channel % 4 != 0 || end_channel % 4 != 0 float maximum {
int slice_coord_TMP = (start_channel) / 4 int sub_ch_coord_TMP = (start_channel) % 4 float4 src_TMP = convert_float4(
read_imageh(src_tensor_image_buffer, ((slice_coord_TMP) *shared_int4_1.x + (Y)) * shared_int4_1.y + (X))) maximum = (float[4]){ src_TMP.x, src_TMP.y, src_TMP.z, src_TMP.w }[sub_ch_coord_TMP] } for (int d = start_slice float4 mask_dot = (float4) (1.f) float4 src =
convert_float4(read_imageh(src_tensor_image_buffer, ((d) *shared_int4_1.x + (Y)) * shared_int4_1.y + (X))) if (need_per_channels_check && (d == start_slice || d == end_slice - 1)) {
if (d * 4 + 0 < start_channel || d * 4 + 0 >= end_channel) {
mask_dot.x = 0.f src.x = maximum }
if (d * 4 + 1 < start_channel || d * 4 + 1 >= end_channel) {
mask_dot.y = 0.f src.y = maximum }
if (d * 4 + 2 < start_channel || d * 4 + 2 >= end_channel) {
mask_dot.z = 0.f src.z = maximum }
if (d * 4 + 3 < start_channel || d * 4 + 3 >= end_channel) {
mask_dot.w = 0.f src.w = maximum }
}
float new_max = max(src.x, src.y) new_max = max(new_max, src.z) new_max = max(new_max, src.w) new_max = max(new_max, maximum) float scale = native_exp(maximum - new_max) maximum = new_max sum *= scale float4 exp_res = native_exp(src - maximum) sum += dot(mask_dot, exp_res) }
if (!isfinite(maximum) || sum == 0.0f) {
write_imageh(dst_tensor_image2d, (int2) (X, Y), (half4) (0.0h)) return }
write_imageh(dst_tensor_image2d, (int2) (X, Y),
(half4) (convert_half(1.0f / sum), convert_half(maximum), 0.0h, 0.0h))}
__kernel void adreno_xmem_attn_softmax_apply_basic(global half4 * dst_tensor_buffer,
read_only image1d_buffer_t src_tensor_image_buffer,
read_only image2d_t src_tensor_1_image2d,
const int4 shared_int4_0,
const int4 shared_int4_1) {
int X = get_global_id(0) int Y = get_global_id(1) int Z = get_global_id(2) if (X >= shared_int4_0.z || Y >= shared_int4_0.x || Z >= shared_int4_0.y) {
return }
half4 src = read_imageh(src_tensor_image_buffer, ((Z) *shared_int4_1.x + (Y)) * shared_int4_1.y + (X)) {
half4 src_final {
{
half4 exp_val = read_imageh(src_tensor_1_image2d, smp_zero, (int2) (X, Y)) src_final = exp(src - exp_val.y) * exp_val.x const int k = Z * 4 const int n_kv = shared_int4_1.z if (k + 0 >= n_kv) {
src_final.x = 0.0h }
if (k + 1 >= n_kv) {
src_final.y = 0.0h }
if (k + 2 >= n_kv) {
src_final.z = 0.0h }
if (k + 3 >= n_kv) {
src_final.w = 0.0h }
}
}
dst_tensor_buffer[(((Z) *shared_int4_0.x + (Y)) * shared_int4_0.z + (X))] = src_final }}
__kernel void adreno_xmem_attn_mask_scores(global half4 * dst_score_tensor_buffer,
read_only image1d_buffer_t src_score_image_buffer,
const global half * mask,
const ulong mask_offset,
const int q_width,
const int n_q,
const int n_kv,
const int n_kv_padded,
const int kv_heads_total,
const int n_head,
const int n_head_kv,
const ulong mask_nb1,
const ulong mask_nb2,
const ulong mask_nb3,
const int mask_ne2,
const int mask_ne3) {
const int X = get_global_id(0) const int Y = get_global_id(1) const int Z = get_global_id(2) const int npack = n_kv_padded / 4 if (X >= q_width || Y >= kv_heads_total || Z >= npack) {
return }
const int gqa = n_head / n_head_kv const int head_kv = Y % n_head_kv const int batch = Y / n_head_kv const int head_group = X / n_q const int q = X - head_group * n_q const int head = head_kv * gqa + head_group const int mask_head_idx = head % mask_ne2 const int mask_batch_idx = batch % mask_ne3 const global char * mask_base = (const global char *) mask + mask_offset const global half * mask_row = (const global half *) (mask_base + mask_batch_idx * mask_nb3 +
mask_head_idx * mask_nb2 + q * mask_nb1)
const half4 score = read_imageh(src_score_image_buffer, ((Z * kv_heads_total + Y) * q_width + X)) float vals[4] = {
convert_float(score.x),
convert_float(score.y),
convert_float(score.z),
convert_float(score.w),
}
for (int lane = 0 const int k_idx = Z * 4 + lane if (k_idx >= n_kv) {
vals[lane] = -INFINITY } else {
vals[lane] += convert_float(mask_row[k_idx]) }
}
dst_score_tensor_buffer[((Z * kv_heads_total + Y) * q_width + X)] =
(half4) (convert_half(vals[0]), convert_half(vals[1]), convert_half(vals[2]), convert_half(vals[3]))}
__kernel void adreno_xmem_attn_pack_v(global half4 * dst_tensor_buffer,
read_only image2d_t src_image2d,
const int4 shared_int4_0,
const int4 shared_int4_1) {
int linear_index = get_global_id(0) if (linear_index >= shared_int4_0.y) {
return }
if (get_global_id(1) != 0) {
return }
if (get_global_id(2) != 0) {
return }
int dst_o_sp_i_ogroup = linear_index int dst_ogroup = dst_o_sp_i_ogroup % shared_int4_0.x int dst_o_sp_i = dst_o_sp_i_ogroup / shared_int4_0.x int dst_i = dst_o_sp_i % shared_int4_0.z int dst_o_sp = dst_o_sp_i / shared_int4_0.z int dst_sp = dst_o_sp % shared_int4_1.x int dst_o = dst_o_sp / shared_int4_1.x int i_slice = dst_i int o_slice = dst_o * shared_int4_0.x + dst_ogroup int spatial_linear = dst_sp int W = spatial_linear % shared_int4_1.y int H = spatial_linear / shared_int4_1.y half4 w0 = (half4) (0) half4 w1 = (half4) (0) half4 w2 = (half4) (0) half4 w3 = (half4) (0)
if (i_slice * 4 < shared_int4_0.w && o_slice < shared_int4_1.z) {
w0 = read_imageh(src_image2d, smp_zero, (int2) ((i_slice * 4), ((W) *shared_int4_1.z + (o_slice)))) }
if (i_slice * 4 + 1 < shared_int4_0.w && o_slice < shared_int4_1.z) {
w1 = read_imageh(src_image2d, smp_zero, (int2) ((i_slice * 4 + 1), ((W) *shared_int4_1.z + (o_slice)))) }
if (i_slice * 4 + 2 < shared_int4_0.w && o_slice < shared_int4_1.z) {
w2 = read_imageh(src_image2d, smp_zero, (int2) ((i_slice * 4 + 2), ((W) *shared_int4_1.z + (o_slice)))) }
if (i_slice * 4 + 3 < shared_int4_0.w && o_slice < shared_int4_1.z) {
w3 = read_imageh(src_image2d, smp_zero, (int2) ((i_slice * 4 + 3), ((W) *shared_int4_1.z + (o_slice)))) }
half4 r0 = w0 half4 r1 = w1 half4 r2 = w2 half4 r3 = w3 dst_tensor_buffer[linear_index * 4 + 0] = r0 dst_tensor_buffer[linear_index * 4 + 1] = r1 dst_tensor_buffer[linear_index * 4 + 2] = r2 dst_tensor_buffer[linear_index * 4 + 3] = r3}
__attribute__((qcom_max_concurrent_subgroups(12))) __kernel void adreno_xmem_attn_pv_gemm(
constant half8 * weights_buffer __attribute__((sub_group_uniform)),
constant half8 * xmem_buffer __attribute__((max_constant_size((6144)))),
read_only image1d_buffer_t src_tensor_image_buffer,
write_only image2d_t dst_tensor_image2d,
const int4 shared_int4_0,
const int4 shared_int4_1,
const int4 shared_int4_2,
const int4 shared_int4_3) {
int X = get_group_id(1) * get_local_size(0) + get_local_id(0) int Y = get_group_id(2) * get_local_size(1) + get_local_id(1) int Z = get_group_id(0) * get_local_size(2) + get_local_id(2) if (X >= shared_int4_0.z || Y >= shared_int4_0.x) {
return }
if (Z * 8 >= shared_int4_0.y) {
return }
half4 r0 = (half4) (0.f) half4 r1 = (half4) (0.f) half4 r2 = (half4) (0.f) half4 r3 = (half4) (0.f) half4 r4 = (half4) (0.f) half4 r5 = (half4) (0.f) half4 r6 = (half4) (0.f) half4 r7 = (half4) (0.f) int x_coord = mad24(X, shared_int4_2.w, shared_int4_1.y) int y_coord = mad24(Y, shared_int4_3.x, shared_int4_1.z) int coord_x, coord_y, coord_s int f_offset = (Z * shared_int4_1.w + Y) * shared_int4_1.x * 32
int subgroup_id = (int) ((0x1F & qcom_get_physical_sub_group_id())) subgroup_id = subgroup_id % 12 int c_offset = mul24(subgroup_id, shared_int4_0.w) __constant half16 * weights_cache = (__constant half16 *) &xmem_buffer[c_offset] coord_y = Y coord_x = X int addr = (((0) * shared_int4_1.w + (coord_y)) * shared_int4_2.z + (coord_x)) int dz = shared_int4_2.x coord_s = 0 do {
half4 src0 = read_imageh(src_tensor_image_buffer, addr) addr += dz coord_s++ half4 src1 = read_imageh(src_tensor_image_buffer, addr) addr += dz coord_s++ qcom_sub_group_constant_load8(xmem_buffer, weights_buffer, c_offset, f_offset >> 1, 32) f_offset += 64 qcom_sub_group_sync(QCOM_CLK_CONST_LOAD_SYNC) r0 += src0.x * weights_cache[0].s0123 r0 += src0.y * weights_cache[0].s4567 r0 += src0.z * weights_cache[0].s89ab r0 += src0.w * weights_cache[0].scdef r1 += src0.x * weights_cache[1].s0123 r1 += src0.y * weights_cache[1].s4567 r1 += src0.z * weights_cache[1].s89ab r1 += src0.w * weights_cache[1].scdef r2 += src0.x * weights_cache[2].s0123 r2 += src0.y * weights_cache[2].s4567 r2 += src0.z * weights_cache[2].s89ab r2 += src0.w * weights_cache[2].scdef r3 += src0.x * weights_cache[3].s0123 r3 += src0.y * weights_cache[3].s4567 r3 += src0.z * weights_cache[3].s89ab r3 += src0.w * weights_cache[3].scdef r4 += src0.x * weights_cache[4].s0123 r4 += src0.y * weights_cache[4].s4567 r4 += src0.z * weights_cache[4].s89ab r4 += src0.w * weights_cache[4].scdef r5 += src0.x * weights_cache[5].s0123 r5 += src0.y * weights_cache[5].s4567 r5 += src0.z * weights_cache[5].s89ab r5 += src0.w * weights_cache[5].scdef r6 += src0.x * weights_cache[6].s0123 r6 += src0.y * weights_cache[6].s4567 r6 += src0.z * weights_cache[6].s89ab r6 += src0.w * weights_cache[6].scdef r7 += src0.x * weights_cache[7].s0123 r7 += src0.y * weights_cache[7].s4567 r7 += src0.z * weights_cache[7].s89ab r7 += src0.w * weights_cache[7].scdef r0 += src1.x * weights_cache[8].s0123 r0 += src1.y * weights_cache[8].s4567 r0 += src1.z * weights_cache[8].s89ab r0 += src1.w * weights_cache[8].scdef r1 += src1.x * weights_cache[9].s0123 r1 += src1.y * weights_cache[9].s4567 r1 += src1.z * weights_cache[9].s89ab r1 += src1.w * weights_cache[9].scdef r2 += src1.x * weights_cache[10].s0123 r2 += src1.y * weights_cache[10].s4567 r2 += src1.z * weights_cache[10].s89ab r2 += src1.w * weights_cache[10].scdef r3 += src1.x * weights_cache[11].s0123 r3 += src1.y * weights_cache[11].s4567 r3 += src1.z * weights_cache[11].s89ab r3 += src1.w * weights_cache[11].scdef r4 += src1.x * weights_cache[12].s0123 r4 += src1.y * weights_cache[12].s4567 r4 += src1.z * weights_cache[12].s89ab r4 += src1.w * weights_cache[12].scdef r5 += src1.x * weights_cache[13].s0123 r5 += src1.y * weights_cache[13].s4567 r5 += src1.z * weights_cache[13].s89ab r5 += src1.w * weights_cache[13].scdef r6 += src1.x * weights_cache[14].s0123 r6 += src1.y * weights_cache[14].s4567 r6 += src1.z * weights_cache[14].s89ab r6 += src1.w * weights_cache[14].scdef r7 += src1.x * weights_cache[15].s0123 r7 += src1.y * weights_cache[15].s4567 r7 += src1.z * weights_cache[15].s89ab r7 += src1.w * weights_cache[15].scdef } while (coord_s < shared_int4_2.y)
coord_s = mul24(Z, 8) coord_x = X coord_y = Y if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r0) if (coord_s < 0) {
res += read_imageh(src_tensor_image_buffer, ((0) * shared_int4_1.w + (0)) * shared_int4_2.z + (0)) }
write_imageh(dst_tensor_image2d, (int2) ((coord_x), ((coord_y) *shared_int4_0.y + (coord_s))), res) coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r1) if (coord_s < 0) {
res += read_imageh(src_tensor_image_buffer, ((0) * shared_int4_1.w + (0)) * shared_int4_2.z + (0)) }
write_imageh(dst_tensor_image2d, (int2) ((coord_x), ((coord_y) *shared_int4_0.y + (coord_s))), res) coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r2) if (coord_s < 0) {
res += read_imageh(src_tensor_image_buffer, ((0) * shared_int4_1.w + (0)) * shared_int4_2.z + (0)) }
write_imageh(dst_tensor_image2d, (int2) ((coord_x), ((coord_y) *shared_int4_0.y + (coord_s))), res) coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r3) if (coord_s < 0) {
res += read_imageh(src_tensor_image_buffer, ((0) * shared_int4_1.w + (0)) * shared_int4_2.z + (0)) }
write_imageh(dst_tensor_image2d, (int2) ((coord_x), ((coord_y) *shared_int4_0.y + (coord_s))), res) coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r4) if (coord_s < 0) {
res += read_imageh(src_tensor_image_buffer, ((0) * shared_int4_1.w + (0)) * shared_int4_2.z + (0)) }
write_imageh(dst_tensor_image2d, (int2) ((coord_x), ((coord_y) *shared_int4_0.y + (coord_s))), res) coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r5) if (coord_s < 0) {
res += read_imageh(src_tensor_image_buffer, ((0) * shared_int4_1.w + (0)) * shared_int4_2.z + (0)) }
write_imageh(dst_tensor_image2d, (int2) ((coord_x), ((coord_y) *shared_int4_0.y + (coord_s))), res) coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r6) if (coord_s < 0) {
res += read_imageh(src_tensor_image_buffer, ((0) * shared_int4_1.w + (0)) * shared_int4_2.z + (0)) }
write_imageh(dst_tensor_image2d, (int2) ((coord_x), ((coord_y) *shared_int4_0.y + (coord_s))), res) coord_s++ }
if (coord_s < shared_int4_0.y) {
half4 res = convert_half4(r7) if (coord_s < 0) {
res += read_imageh(src_tensor_image_buffer, ((0) * shared_int4_1.w + (0)) * shared_int4_2.z + (0)) }
write_imageh(dst_tensor_image2d, (int2) ((coord_x), ((coord_y) *shared_int4_0.y + (coord_s))), res) coord_s++ }
}