diff --git a/ggml/src/ggml-opencl/CMakeLists.txt b/ggml/src/ggml-opencl/CMakeLists.txt index 8a1b6b964..8e6e374f3 100644 --- a/ggml/src/ggml-opencl/CMakeLists.txt +++ b/ggml/src/ggml-opencl/CMakeLists.txt @@ -238,7 +238,7 @@ set(GGML_OPENCL_KERNELS ) if (GGML_OPENCL_USE_ADRENO_KERNELS) - list(APPEND GGML_OPENCL_KERNELS gemm_xmem_f16_f32_os8) + list(APPEND GGML_OPENCL_KERNELS gemm_xmem_f16_f32_os8 sdpa_xmem_f32_f16_os8) endif () foreach (K ${GGML_OPENCL_KERNELS}) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 12465a517..864da2536 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -417,6 +417,10 @@ static void populateProfilingInfo( struct ggml_backend_opencl_context; +#ifdef GGML_OPENCL_USE_ADRENO_KERNELS +static void ggml_cl_adreno_xmem_attn_release_scratch(ggml_backend_opencl_context * backend_ctx); +#endif + // backend device context struct ggml_backend_opencl_device_context { cl_platform_id platform; @@ -537,6 +541,54 @@ struct ggml_opencl_fa_kernels { std::set>> variant_attempted; }; +#ifdef GGML_OPENCL_USE_ADRENO_KERNELS +struct ggml_cl_adreno_xmem_attn_scratch { + cl_mem q_img = nullptr; + cl_mem k_img = nullptr; + cl_mem v_img = nullptr; + cl_mem out_img = nullptr; + cl_mem k_transpose_buf = nullptr; + cl_mem k_transpose_img1d = nullptr; + cl_mem k_packed_buf = nullptr; + cl_mem v_packed_buf = nullptr; + cl_mem score_buf = nullptr; + cl_mem prob_buf = nullptr; + cl_mem score_img1d = nullptr; + cl_mem prob_img1d = nullptr; + cl_mem softmax_stats_img2d = nullptr; + cl_mem xmem_qk = nullptr; + cl_mem xmem_pv = nullptr; + + int n_q = 0; + int n_kv = 0; + int n_kv_padded = 0; + int d_head_q = 0; + int d_head_v = 0; + int q_width = 0; + int kv_heads_total = 0; +}; + +struct ggml_cl_adreno_xmem_attn_state { + bool compiled = false; + bool logged = false; + + cl_kernel kernel_q_f32_to_img_scaled = nullptr; + cl_kernel kernel_kv_f32_to_img_gqa = nullptr; + cl_kernel kernel_kv_f16_to_img_gqa = nullptr; + cl_kernel kernel_img_to_f32 = nullptr; + cl_kernel kernel_k_gather = nullptr; + cl_kernel kernel_pack_k = nullptr; + cl_kernel kernel_qk_gemm = nullptr; + cl_kernel kernel_softmax_reduce_basic = nullptr; + cl_kernel kernel_softmax_apply_basic = nullptr; + cl_kernel kernel_mask_scores = nullptr; + cl_kernel kernel_pack_v = nullptr; + cl_kernel kernel_pv_gemm = nullptr; + + ggml_cl_adreno_xmem_attn_scratch scratch; +}; +#endif + // backend context struct ggml_backend_opencl_context { int ref_count; @@ -762,6 +814,9 @@ struct ggml_backend_opencl_context { cl_kernel kernel_soft_max, kernel_soft_max_4; cl_kernel kernel_soft_max_f16, kernel_soft_max_4_f16; ggml_opencl_fa_kernels fa; +#ifdef GGML_OPENCL_USE_ADRENO_KERNELS + ggml_cl_adreno_xmem_attn_state adreno_xmem_attn; +#endif cl_kernel kernel_get_rows_f32, kernel_get_rows_f16, kernel_get_rows_q4_0; cl_kernel kernel_set_rows_f32_i64, kernel_set_rows_f32_i32, kernel_set_rows_f16_i64, kernel_set_rows_f16_i32; cl_kernel kernel_set_rows_q8_0_i64, kernel_set_rows_q8_0_i32; @@ -1176,6 +1231,9 @@ struct ggml_backend_opencl_context { if (kv.second.image) { CL_CHECK(clReleaseMemObject(kv.second.image)); } } dequant_f16_pool.clear(); +#ifdef GGML_OPENCL_USE_ADRENO_KERNELS + ggml_cl_adreno_xmem_attn_release_scratch(this); +#endif } } }; @@ -2343,6 +2401,49 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { } #endif // GGML_OPENCL_USE_ADRENO_KERNELS +#ifdef GGML_OPENCL_USE_ADRENO_KERNELS + // Adreno xmem SDPA + if (backend_ctx->gpu_family == GPU_FAMILY::ADRENO) { +#ifdef GGML_OPENCL_EMBED_KERNELS + const std::string kernel_src { + #include "sdpa_xmem_f32_f16_os8.cl.h" + }; +#else + const std::string kernel_src = read_file("sdpa_xmem_f32_f16_os8.cl"); +#endif + cl_program program = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts); + + auto & xmem_attn = backend_ctx->adreno_xmem_attn; + CL_CHECK((xmem_attn.kernel_q_f32_to_img_scaled = + clCreateKernel(program, "adreno_xmem_attn_q_f32_to_img_scaled", &err), err)); + CL_CHECK((xmem_attn.kernel_kv_f32_to_img_gqa = + clCreateKernel(program, "adreno_xmem_attn_kv_f32_to_img_gqa", &err), err)); + CL_CHECK((xmem_attn.kernel_kv_f16_to_img_gqa = + clCreateKernel(program, "adreno_xmem_attn_kv_f16_to_img_gqa", &err), err)); + CL_CHECK((xmem_attn.kernel_img_to_f32 = + clCreateKernel(program, "adreno_xmem_attn_img_to_f32", &err), err)); + CL_CHECK((xmem_attn.kernel_k_gather = + clCreateKernel(program, "adreno_xmem_attn_k_gather", &err), err)); + CL_CHECK((xmem_attn.kernel_pack_k = + clCreateKernel(program, "adreno_xmem_attn_pack_k", &err), err)); + CL_CHECK((xmem_attn.kernel_qk_gemm = + clCreateKernel(program, "adreno_xmem_attn_qk_gemm", &err), err)); + CL_CHECK((xmem_attn.kernel_softmax_reduce_basic = + clCreateKernel(program, "adreno_xmem_attn_softmax_reduce_basic", &err), err)); + CL_CHECK((xmem_attn.kernel_softmax_apply_basic = + clCreateKernel(program, "adreno_xmem_attn_softmax_apply_basic", &err), err)); + CL_CHECK((xmem_attn.kernel_mask_scores = + clCreateKernel(program, "adreno_xmem_attn_mask_scores", &err), err)); + CL_CHECK((xmem_attn.kernel_pack_v = + clCreateKernel(program, "adreno_xmem_attn_pack_v", &err), err)); + CL_CHECK((xmem_attn.kernel_pv_gemm = + clCreateKernel(program, "adreno_xmem_attn_pv_gemm", &err), err)); + CL_CHECK(clReleaseProgram(program)); + xmem_attn.compiled = true; + GGML_LOG_CONT("."); + } +#endif // GGML_OPENCL_USE_ADRENO_KERNELS + // mul_mm_f32_f32_l4_lm { #ifdef GGML_OPENCL_EMBED_KERNELS @@ -15927,6 +16028,581 @@ static constexpr int FD_MAX_N_Q_MULTI = 8; static constexpr int FD_MQ_KV_PER_SPLIT = 256; static constexpr int FD_MQ_MAX_SPLITS = 128; +#ifdef GGML_OPENCL_USE_ADRENO_KERNELS +struct ggml_cl_adreno_xmem_attn_schedule { + int qk_lws0 = 256; + int qk_lws2 = 1; + int softmax_reduce_lws0 = 256; + int softmax_apply_lws0 = 64; + int softmax_apply_lws2 = 4; + int pv_lws0 = 64; + int pv_lws2 = 4; +}; + +static inline size_t ggml_cl_round_up(size_t x, size_t a) { + return ((x + a - 1) / a) * a; +} + +static inline int ggml_cl_round_up_div(int x, int y) { + return (x + y - 1) / y; +} + +static inline void ggml_cl_set_arg_int4(cl_kernel kernel, cl_uint index, int x, int y, int z, int w) { + struct { int x, y, z, w; } value { x, y, z, w }; + CL_CHECK(clSetKernelArg(kernel, index, sizeof(value), &value)); +} + +static cl_mem ggml_cl_make_image2d_half4(cl_context context, cl_mem_flags flags, size_t width, size_t height) { + cl_int err = CL_SUCCESS; + cl_image_format format = { CL_RGBA, CL_HALF_FLOAT }; + cl_image_desc desc = {}; + desc.image_type = CL_MEM_OBJECT_IMAGE2D; + desc.image_width = width; + desc.image_height = height; + cl_mem image = clCreateImage(context, flags, &format, &desc, nullptr, &err); + CL_CHECK(err); + return image; +} + +static cl_mem ggml_cl_make_image1d_buffer_half4(cl_context context, cl_mem_flags flags, size_t width, cl_mem backing_buffer) { + cl_int err = CL_SUCCESS; + cl_image_format format = { CL_RGBA, CL_HALF_FLOAT }; + cl_image_desc desc = {}; + desc.image_type = CL_MEM_OBJECT_IMAGE1D_BUFFER; + desc.image_width = width; + desc.buffer = backing_buffer; + cl_mem image = clCreateImage(context, flags, &format, &desc, nullptr, &err); + CL_CHECK(err); + return image; +} + +static void ggml_cl_release_mem(cl_mem & mem) { + if (mem != nullptr) { + CL_CHECK(clReleaseMemObject(mem)); + mem = nullptr; + } +} + +static void ggml_cl_adreno_xmem_attn_release_scratch(ggml_backend_opencl_context * backend_ctx) { + auto & s = backend_ctx->adreno_xmem_attn.scratch; + ggml_cl_release_mem(s.q_img); + ggml_cl_release_mem(s.k_img); + ggml_cl_release_mem(s.v_img); + ggml_cl_release_mem(s.out_img); + ggml_cl_release_mem(s.k_transpose_img1d); + ggml_cl_release_mem(s.k_transpose_buf); + ggml_cl_release_mem(s.k_packed_buf); + ggml_cl_release_mem(s.v_packed_buf); + ggml_cl_release_mem(s.score_img1d); + ggml_cl_release_mem(s.prob_img1d); + ggml_cl_release_mem(s.score_buf); + ggml_cl_release_mem(s.prob_buf); + ggml_cl_release_mem(s.softmax_stats_img2d); + ggml_cl_release_mem(s.xmem_qk); + ggml_cl_release_mem(s.xmem_pv); + s = {}; +} + +static ggml_cl_adreno_xmem_attn_schedule ggml_cl_adreno_xmem_attn_select_schedule( + const ggml_backend_opencl_context * backend_ctx, + int n_q, + int n_kv, + int heads_total, + int q_width, + int gqa_ratio) { + const bool big_h = heads_total >= 8; + ggml_cl_adreno_xmem_attn_schedule sched; + + if (gqa_ratio == 1) { + if (n_q >= 512) { sched.qk_lws0 = 512; } + else if (n_q >= 256) { sched.qk_lws0 = 128; } + else { sched.qk_lws0 = 64; } + sched.qk_lws2 = (big_h && n_q >= 512) ? 2 : 1; + } else { + if (q_width >= 2048) { sched.qk_lws0 = 512; } + else if (q_width >= 256) { sched.qk_lws0 = 128; } + else { sched.qk_lws0 = 64; } + sched.qk_lws2 = MIN(8, (int) backend_ctx->max_workgroup_size / sched.qk_lws0); + } + + if (n_kv >= 2048) { sched.softmax_reduce_lws0 = 1024; } + else if (n_kv >= 512) { sched.softmax_reduce_lws0 = big_h ? 256 : 512; } + else { sched.softmax_reduce_lws0 = 256; } + + if (n_kv < 256) { sched.softmax_apply_lws0 = 64; } + else { sched.softmax_apply_lws0 = big_h ? 128 : 64; } + sched.softmax_apply_lws2 = n_kv >= 512 ? 8 : 4; + + if (n_q < 256) { sched.pv_lws0 = 64; } + else { sched.pv_lws0 = big_h ? 128 : 64; } + sched.pv_lws2 = big_h ? 8 : (n_q <= 256 ? 8 : 4); + + const int max_wg = (int) backend_ctx->max_workgroup_size; + auto fix = [&](int & l0, int & l2) { + while (l0 * l2 > max_wg) { + if (l2 > 1) { l2 /= 2; } + else if (l0 > 32) { l0 /= 2; } + else { break; } + } + }; + fix(sched.qk_lws0, sched.qk_lws2); + fix(sched.softmax_apply_lws0, sched.softmax_apply_lws2); + fix(sched.pv_lws0, sched.pv_lws2); + while (sched.softmax_reduce_lws0 > max_wg) { + sched.softmax_reduce_lws0 /= 2; + } + + return sched; +} + +static bool ggml_cl_adreno_xmem_attn_prepare( + ggml_backend_opencl_context * backend_ctx, + int n_q, + int n_kv, + int d_head_q, + int d_head_v, + int n_head, + int n_head_kv, + int n_batch) { + auto & s = backend_ctx->adreno_xmem_attn.scratch; + const int gqa_ratio = n_head / n_head_kv; + const int q_width = n_q * gqa_ratio; + const int kv_heads_total = n_head_kv * n_batch; + const int n_kv_padded = (int) ggml_cl_round_up((size_t) n_kv, 32); + if (s.q_img != nullptr && + s.n_q == n_q && + s.n_kv == n_kv && + s.n_kv_padded == n_kv_padded && + s.d_head_q == d_head_q && + s.d_head_v == d_head_v && + s.q_width == q_width && + s.kv_heads_total == kv_heads_total) { + return true; + } + + ggml_cl_adreno_xmem_attn_release_scratch(backend_ctx); + + const int qpack = d_head_q / 4; + const int vpack = d_head_v / 4; + const int npack = n_kv_padded / 4; + const size_t q_img_h = (size_t) kv_heads_total * qpack; + const size_t v_img_h = (size_t) kv_heads_total * vpack; + + s.q_img = ggml_cl_make_image2d_half4(backend_ctx->context, CL_MEM_READ_WRITE, (size_t) q_width, q_img_h); + s.k_img = ggml_cl_make_image2d_half4(backend_ctx->context, CL_MEM_READ_WRITE, (size_t) n_kv_padded, q_img_h); + s.v_img = ggml_cl_make_image2d_half4(backend_ctx->context, CL_MEM_READ_WRITE, (size_t) n_kv_padded, v_img_h); + s.out_img = ggml_cl_make_image2d_half4(backend_ctx->context, CL_MEM_READ_WRITE, (size_t) q_width, v_img_h); + + const size_t k_transpose_half4_elems = (size_t) npack * kv_heads_total * d_head_q; + s.k_transpose_buf = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, k_transpose_half4_elems * sizeof(uint16_t) * 4, nullptr, nullptr); + GGML_ASSERT(s.k_transpose_buf != nullptr); + s.k_transpose_img1d = ggml_cl_make_image1d_buffer_half4(backend_ctx->context, CL_MEM_READ_ONLY, k_transpose_half4_elems, s.k_transpose_buf); + + const size_t k_groups16 = (size_t) ggml_cl_round_up_div(kv_heads_total * d_head_q, 16); + const size_t v_groups16 = (size_t) ggml_cl_round_up_div(kv_heads_total * d_head_v, 16); + const size_t k_packed_half4_elems = (size_t) n_kv_padded * k_groups16 * 4; + const size_t v_packed_half4_elems = (size_t) n_kv_padded * v_groups16 * 4; + s.k_packed_buf = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, k_packed_half4_elems * sizeof(uint16_t) * 4, nullptr, nullptr); + s.v_packed_buf = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, v_packed_half4_elems * sizeof(uint16_t) * 4, nullptr, nullptr); + GGML_ASSERT(s.k_packed_buf != nullptr && s.v_packed_buf != nullptr); + + const size_t score_half4_elems = (size_t) npack * kv_heads_total * q_width; + const size_t score_bytes = score_half4_elems * sizeof(uint16_t) * 4; + s.score_buf = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, score_bytes, nullptr, nullptr); + s.prob_buf = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, score_bytes, nullptr, nullptr); + GGML_ASSERT(s.score_buf != nullptr && s.prob_buf != nullptr); + s.score_img1d = ggml_cl_make_image1d_buffer_half4(backend_ctx->context, CL_MEM_READ_ONLY, score_half4_elems, s.score_buf); + s.prob_img1d = ggml_cl_make_image1d_buffer_half4(backend_ctx->context, CL_MEM_READ_ONLY, score_half4_elems, s.prob_buf); + s.softmax_stats_img2d = ggml_cl_make_image2d_half4(backend_ctx->context, CL_MEM_READ_WRITE, + (size_t) q_width, (size_t) kv_heads_total); + s.xmem_qk = clCreateBuffer(backend_ctx->context, CL_MEM_READ_ONLY, 6144, nullptr, nullptr); + s.xmem_pv = clCreateBuffer(backend_ctx->context, CL_MEM_READ_ONLY, 6144, nullptr, nullptr); + GGML_ASSERT(s.softmax_stats_img2d != nullptr && s.xmem_qk != nullptr && s.xmem_pv != nullptr); + + s.n_q = n_q; + s.n_kv = n_kv; + s.n_kv_padded = n_kv_padded; + s.d_head_q = d_head_q; + s.d_head_v = d_head_v; + s.q_width = q_width; + s.kv_heads_total = kv_heads_total; + return true; +} + +static bool ggml_cl_adreno_xmem_attn_can_use( + const ggml_backend_opencl_context * backend_ctx, + const ggml_tensor * q, + const ggml_tensor * k, + const ggml_tensor * dst) { + static const char * xmem_sdpa_env = getenv("GGML_OPENCL_XMEM_SDPA"); + if (xmem_sdpa_env == nullptr || xmem_sdpa_env[0] == '0') { + return false; + } + + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; + + if (!backend_ctx->adreno_xmem_attn.compiled || backend_ctx->gpu_family != GPU_FAMILY::ADRENO) { + return false; + } + if (q->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + (k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_F32) || + (v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_F32)) { + return false; + } + if (sinks != nullptr) { + return false; + } + if (q->nb[0] != ggml_type_size(q->type) || k->nb[0] != ggml_type_size(k->type) || + v->nb[0] != ggml_type_size(v->type) || dst->nb[0] != ggml_type_size(dst->type)) { + return false; + } + if (mask != nullptr && (mask->type != GGML_TYPE_F16 || mask->nb[0] != sizeof(ggml_fp16_t))) { + return false; + } + + const int n_q = q->ne[1]; + const int n_kv = k->ne[1]; + const int d_head_q = q->ne[0]; + const int d_head_v = v->ne[0]; + const int n_head = q->ne[2]; + const int n_head_kv = k->ne[2]; + const int n_batch = q->ne[3]; + + if (n_q <= 1 || n_kv <= 0 || n_kv > 8192) { + return false; + } + if (d_head_q != k->ne[0] || d_head_v != v->ne[0] || k->ne[1] != v->ne[1] || k->ne[3] != v->ne[3]) { + return false; + } + if (q->ne[3] != k->ne[3]) { + return false; + } + if (n_head_kv <= 0 || n_head % n_head_kv != 0 || k->ne[2] != v->ne[2]) { + return false; + } + if (dst->ne[0] != d_head_v || dst->ne[1] != n_head || dst->ne[2] != n_q || dst->ne[3] != n_batch) { + return false; + } + if ((d_head_q % 8) != 0 || (d_head_v % 32) != 0) { + return false; + } + if (mask != nullptr && + (mask->ne[0] < n_kv || mask->ne[1] < n_q || mask->ne[2] <= 0 || mask->ne[3] <= 0)) { + return false; + } + + float params[3]; + memcpy(params, dst->op_params, sizeof(params)); + if (params[1] != 0.0f || params[2] != 0.0f) { + return false; + } + + const int gqa_ratio = n_head / n_head_kv; + const int q_width = n_q * gqa_ratio; + const int kv_heads_total = n_head_kv * n_batch; + const int n_kv_padded = (int) ggml_cl_round_up((size_t) n_kv, 32); + const int qpack = d_head_q / 4; + const int vpack = d_head_v / 4; + const int npack = n_kv_padded / 4; + + if ((size_t) q_width > backend_ctx->image2d_max_width || + (size_t) n_kv_padded > backend_ctx->image2d_max_width) { + return false; + } + if ((size_t) kv_heads_total * (size_t) qpack > backend_ctx->image2d_max_height || + (size_t) kv_heads_total * (size_t) vpack > backend_ctx->image2d_max_height) { + return false; + } + if ((size_t) npack * (size_t) kv_heads_total * (size_t) d_head_q > backend_ctx->image_max_buffer_size || + (size_t) npack * (size_t) kv_heads_total * (size_t) q_width > backend_ctx->image_max_buffer_size) { + return false; + } + + return true; +} + +static void ggml_cl_adreno_xmem_attn_run( + ggml_backend_t backend, + const ggml_tensor * q, + const ggml_tensor * k, + ggml_tensor * dst) { + ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *) backend->context; + auto & xstate = backend_ctx->adreno_xmem_attn; + auto & s = xstate.scratch; + if (!xstate.logged) { + GGML_LOG_INFO("ggml_opencl: using Adreno xmem attention path\n"); + xstate.logged = true; + } + + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + + ggml_tensor_extra_cl * extra_q = (ggml_tensor_extra_cl *) q->extra; + ggml_tensor_extra_cl * extra_k = (ggml_tensor_extra_cl *) k->extra; + ggml_tensor_extra_cl * extra_v = (ggml_tensor_extra_cl *) v->extra; + ggml_tensor_extra_cl * extra_o = (ggml_tensor_extra_cl *) dst->extra; + ggml_tensor_extra_cl * extra_mask = mask ? (ggml_tensor_extra_cl *) mask->extra : nullptr; + + const cl_ulong offset_q = extra_q->offset + q->view_offs; + const cl_ulong offset_k = extra_k->offset + k->view_offs; + const cl_ulong offset_v = extra_v->offset + v->view_offs; + const cl_ulong offset_o = extra_o->offset + dst->view_offs; + const cl_ulong offset_mask = extra_mask ? extra_mask->offset + mask->view_offs : 0; + + const int n_q = q->ne[1]; + const int n_kv = k->ne[1]; + const int d_head_q = q->ne[0]; + const int d_head_v = v->ne[0]; + const int n_head = q->ne[2]; + const int n_head_kv = k->ne[2]; + const int n_batch = q->ne[3]; + const int heads_total = n_head * n_batch; + const int gqa_ratio = n_head / n_head_kv; + const int q_width = n_q * gqa_ratio; + const int kv_heads_total = n_head_kv * n_batch; + const int n_kv_padded = (int) ggml_cl_round_up((size_t) n_kv, 32); + const int qpack = d_head_q / 4; + const int opack = d_head_v / 4; + const int npack = n_kv_padded / 4; + const float scale = ((const float *) dst->op_params)[0]; + + GGML_ASSERT(ggml_cl_adreno_xmem_attn_prepare( + backend_ctx, n_q, n_kv, d_head_q, d_head_v, n_head, n_head_kv, n_batch)); + const ggml_cl_adreno_xmem_attn_schedule sched = + ggml_cl_adreno_xmem_attn_select_schedule( + backend_ctx, n_q, n_kv_padded, heads_total, q_width, gqa_ratio); + + { + size_t gws[3] = {ggml_cl_round_up((size_t) n_q, 8), (size_t) heads_total, (size_t) qpack}; + size_t lws[3] = {8, 1, (size_t) ((qpack <= 32) ? qpack : 1)}; + cl_kernel kernel = xstate.kernel_q_f32_to_img_scaled; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra_q->data_device)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset_q)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &s.q_img)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(float), &scale)); + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &d_head_q)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &n_q)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &n_head)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(int), &n_head_kv)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(int), &n_batch)); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &q->nb[1])); + CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_ulong), &q->nb[2])); + CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &q->nb[3])); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + + { + size_t gws[3] = {(size_t) n_kv_padded, (size_t) kv_heads_total, (size_t) qpack}; + size_t lws[3] = {8, 1, (size_t) ((qpack <= 32) ? qpack : 1)}; + cl_kernel kernel = k->type == GGML_TYPE_F16 ? + xstate.kernel_kv_f16_to_img_gqa : xstate.kernel_kv_f32_to_img_gqa; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra_k->data_device)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset_k)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &s.k_img)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(int), &d_head_q)); + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &n_kv)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &n_kv_padded)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &n_head_kv)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(int), &n_batch)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_ulong), &k->nb[1])); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &k->nb[2])); + CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_ulong), &k->nb[3])); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + + { + size_t gws[3] = {(size_t) n_kv_padded, (size_t) kv_heads_total, (size_t) opack}; + size_t lws[3] = {8, 1, (size_t) ((opack <= 32) ? opack : 1)}; + cl_kernel kernel = v->type == GGML_TYPE_F16 ? + xstate.kernel_kv_f16_to_img_gqa : xstate.kernel_kv_f32_to_img_gqa; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra_v->data_device)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset_v)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &s.v_img)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(int), &d_head_v)); + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &n_kv)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &n_kv_padded)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &n_head_kv)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(int), &n_batch)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_ulong), &v->nb[1])); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &v->nb[2])); + CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_ulong), &v->nb[3])); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + + { + size_t gws[3] = {(size_t) d_head_q, (size_t) kv_heads_total, (size_t) npack}; + size_t lws[3] = {(size_t) MIN(64, d_head_q), (size_t) (kv_heads_total >= 2 ? 2 : 1), (size_t) MIN(8, npack)}; + if (lws[0] * lws[1] * lws[2] > backend_ctx->max_workgroup_size) { + lws[1] = 1; + } + cl_kernel kernel = xstate.kernel_k_gather; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &s.k_transpose_buf)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &s.k_img)); + ggml_cl_set_arg_int4(kernel, 2, n_kv_padded, kv_heads_total, npack, d_head_q); + ggml_cl_set_arg_int4(kernel, 3, qpack, 0, 0, 0); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + { + const size_t groups16 = (size_t) ggml_cl_round_up_div(kv_heads_total * d_head_q, 16); + const size_t packed_linear = (size_t) n_kv_padded * groups16; + const size_t lws0 = MIN((size_t) 1024, backend_ctx->max_workgroup_size); + size_t gws[3] = {ggml_cl_round_up(packed_linear, lws0), 1, 1}; + size_t lws[3] = {lws0, 1, 1}; + cl_kernel kernel = xstate.kernel_pack_k; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &s.k_packed_buf)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &s.k_transpose_img1d)); + ggml_cl_set_arg_int4(kernel, 2, 8, (int) packed_linear, qpack, d_head_q); + ggml_cl_set_arg_int4(kernel, 3, kv_heads_total, kv_heads_total, kv_heads_total, npack); + ggml_cl_set_arg_int4(kernel, 4, d_head_q, 0, 0, 0); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + + { + size_t lws[3] = {(size_t) sched.qk_lws0, 1, (size_t) sched.qk_lws2}; + const int slices_per_group = sched.qk_lws2 * 8; + const size_t groups_z = (size_t) ggml_cl_round_up_div(npack, slices_per_group); + const size_t groups_x = (size_t) ggml_cl_round_up_div(q_width, sched.qk_lws0); + size_t gws[3] = { + lws[0] * groups_z, + groups_x, + (size_t) kv_heads_total * lws[2], + }; + + cl_kernel kernel = xstate.kernel_qk_gemm; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &s.score_buf)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &s.k_packed_buf)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &s.xmem_qk)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem), &s.q_img)); + ggml_cl_set_arg_int4(kernel, 4, kv_heads_total, npack, q_width, 32); + ggml_cl_set_arg_int4(kernel, 5, qpack, 0, 0, kv_heads_total); + ggml_cl_set_arg_int4(kernel, 6, qpack, 1, 1, 0); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + cl_mem softmax_input_img = s.score_img1d; + cl_mem softmax_output_buf = s.prob_buf; + cl_mem pv_prob_img = s.prob_img1d; + + if (mask != nullptr) { + const cl_ulong mask_nb1 = mask->nb[1]; + const cl_ulong mask_nb2 = mask->nb[2]; + const cl_ulong mask_nb3 = mask->nb[3]; + const int mask_ne2 = mask->ne[2]; + const int mask_ne3 = mask->ne[3]; + size_t lws[3] = {(size_t) sched.softmax_apply_lws0, 1, (size_t) sched.softmax_apply_lws2}; + size_t gws[3] = { + ggml_cl_round_up((size_t) q_width, lws[0]), + (size_t) kv_heads_total, + ggml_cl_round_up((size_t) npack, lws[2]), + }; + cl_kernel kernel = xstate.kernel_mask_scores; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &s.prob_buf)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &s.score_img1d)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extra_mask->data_device)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset_mask)); + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &q_width)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &n_q)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &n_kv)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(int), &n_kv_padded)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(int), &kv_heads_total)); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(int), &n_head)); + CL_CHECK(clSetKernelArg(kernel, 10, sizeof(int), &n_head_kv)); + CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &mask_nb1)); + CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_ulong), &mask_nb2)); + CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &mask_nb3)); + CL_CHECK(clSetKernelArg(kernel, 14, sizeof(int), &mask_ne2)); + CL_CHECK(clSetKernelArg(kernel, 15, sizeof(int), &mask_ne3)); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + + softmax_input_img = s.prob_img1d; + softmax_output_buf = s.score_buf; + pv_prob_img = s.score_img1d; + } + + { + size_t lws[3] = {(size_t) sched.softmax_reduce_lws0, 1, 1}; + size_t gws[3] = {ggml_cl_round_up((size_t) q_width, lws[0]), (size_t) kv_heads_total, 1}; + cl_kernel kernel = xstate.kernel_softmax_reduce_basic; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &softmax_input_img)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &s.softmax_stats_img2d)); + ggml_cl_set_arg_int4(kernel, 2, kv_heads_total, 1, q_width, n_kv); + ggml_cl_set_arg_int4(kernel, 3, kv_heads_total, q_width, 0, 0); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + { + size_t lws[3] = {(size_t) sched.softmax_apply_lws0, 1, (size_t) sched.softmax_apply_lws2}; + size_t gws[3] = { + ggml_cl_round_up((size_t) q_width, lws[0]), + (size_t) kv_heads_total, + ggml_cl_round_up((size_t) npack, lws[2]), + }; + cl_kernel kernel = xstate.kernel_softmax_apply_basic; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &softmax_output_buf)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &softmax_input_img)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &s.softmax_stats_img2d)); + ggml_cl_set_arg_int4(kernel, 3, kv_heads_total, npack, q_width, 1); + ggml_cl_set_arg_int4(kernel, 4, kv_heads_total, q_width, n_kv, 0); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + { + const size_t groups16 = (size_t) ggml_cl_round_up_div(kv_heads_total * d_head_v, 16); + const size_t packed_linear = (size_t) n_kv_padded * groups16; + const size_t lws0 = MIN((size_t) 1024, backend_ctx->max_workgroup_size); + size_t gws[3] = {ggml_cl_round_up(packed_linear, lws0), 1, 1}; + size_t lws[3] = {lws0, 1, 1}; + cl_kernel kernel = xstate.kernel_pack_v; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &s.v_packed_buf)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &s.v_img)); + ggml_cl_set_arg_int4(kernel, 2, 8, (int) packed_linear, npack, n_kv_padded); + ggml_cl_set_arg_int4(kernel, 3, kv_heads_total, kv_heads_total, opack, 0); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + + { + size_t lws[3] = {(size_t) sched.pv_lws0, 1, (size_t) sched.pv_lws2}; + const int blocks = ggml_cl_round_up_div(opack, 8); + const size_t groups_z = (size_t) ggml_cl_round_up_div(blocks, sched.pv_lws2); + const size_t groups_x = (size_t) ggml_cl_round_up_div(q_width, sched.pv_lws0); + size_t gws[3] = { + lws[0] * groups_z, + groups_x, + (size_t) kv_heads_total * lws[2], + }; + + cl_kernel kernel = xstate.kernel_pv_gemm; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &s.v_packed_buf)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &s.xmem_pv)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &pv_prob_img)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_mem), &s.out_img)); + ggml_cl_set_arg_int4(kernel, 4, kv_heads_total, opack, q_width, 32); + ggml_cl_set_arg_int4(kernel, 5, npack, 0, 0, kv_heads_total); + ggml_cl_set_arg_int4(kernel, 6, kv_heads_total * q_width, npack, q_width, 1); + ggml_cl_set_arg_int4(kernel, 7, 1, 0, 0, 0); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } + + { + size_t gws[3] = {ggml_cl_round_up((size_t) n_q, 8), (size_t) heads_total, (size_t) opack}; + size_t lws[3] = {8, 1, (size_t) ((opack <= 32) ? opack : 1)}; + cl_kernel kernel = xstate.kernel_img_to_f32; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra_o->data_device)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset_o)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &s.out_img)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(int), &d_head_v)); + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &n_q)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &n_head)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &n_head_kv)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(int), &n_batch)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_ulong), &dst->nb[1])); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &dst->nb[2])); + CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_ulong), &dst->nb[3])); + backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst); + } +} + +#endif // GGML_OPENCL_USE_ADRENO_KERNELS + static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, const ggml_tensor * k, ggml_tensor * dst) { const ggml_tensor * v = dst->src[2]; const ggml_tensor * mask = dst->src[3]; @@ -15954,6 +16630,13 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co const int n_head_kv = k->ne[2]; const int n_batch = q->ne[3]; +#ifdef GGML_OPENCL_USE_ADRENO_KERNELS + if (ggml_cl_adreno_xmem_attn_can_use(backend_ctx, q, k, dst)) { + ggml_cl_adreno_xmem_attn_run(backend, q, k, dst); + return; + } +#endif + // DK=512 (Gemma-4 global layers) runs decode-only (q1 / q1_split) on // Adreno - it never uses the BM-tile path, and the prepass + split-tile // programs OOM the compiler at DK=512; supports_op only admits diff --git a/ggml/src/ggml-opencl/kernels/sdpa_xmem_f32_f16_os8.cl b/ggml/src/ggml-opencl/kernels/sdpa_xmem_f32_f16_os8.cl new file mode 100644 index 000000000..26f0fbd52 --- /dev/null +++ b/ggml/src/ggml-opencl/kernels/sdpa_xmem_f32_f16_os8.cl @@ -0,0 +1,871 @@ +#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; i < 4; ++i) { + 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; d < end_slice; d += 1) { + 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; lane < 4; ++lane) { + 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++; + } +}