Commit 9a4580b61 for llama.cpp
commit 9a4580b61cba91a1e3fb5ddccc939c6c73454b66
Author: lhez <lih@qti.qualcomm.com>
Date: Sun Oct 25 02:11:15 2026 -0700
opencl: add topk and lightning indexer (#31026)
* opencl: add topk
* opencl: optimize topk
* opencl: add fp lightning indexer
* opencl: optimize lightning indexer
* opencl: refactor
* opencl: refactor tiled topk kernel
* opencl: adjust format
diff --git a/ggml/src/ggml-opencl/CMakeLists.txt b/ggml/src/ggml-opencl/CMakeLists.txt
index b705e4e33..57620e560 100644
--- a/ggml/src/ggml-opencl/CMakeLists.txt
+++ b/ggml/src/ggml-opencl/CMakeLists.txt
@@ -66,6 +66,8 @@ set(GGML_OPENCL_KERNELS
moe_add_id_glu
argmax
argsort
+ top_k
+ lightning_indexer
tri
fill
clamp
diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp
index c441f6443..65a5e6621 100644
--- a/ggml/src/ggml-opencl/ggml-opencl.cpp
+++ b/ggml/src/ggml-opencl/ggml-opencl.cpp
@@ -663,6 +663,7 @@ struct ggml_backend_opencl_context {
size_t max_workgroup_size;
bool fp16_support;
bool has_vector_subgroup_broadcast;
+ bool has_khr_subgroup_shuffle = false;
bool has_subgroup_shuffle = false; // cl_khr_subgroup_shuffle or cl_qcom_subgroup_shuffle
bool has_integer_dot = false; // cl_khr_integer_dot_product or cl_qcom_dot_product8
bool has_qcom_subgroup_shuffle = false; // specifically cl_qcom_subgroup_shuffle
@@ -755,6 +756,7 @@ struct ggml_backend_opencl_context {
// prealloc buffers for src0 and src1
ggml_cl_buffer prealloc_src0;
ggml_cl_buffer prealloc_src1;
+ ggml_cl_buffer prealloc_top_k;
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
ggml_cl_buffer prealloc_adreno_xmem_const;
@@ -869,6 +871,8 @@ struct ggml_backend_opencl_context {
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
ggml_cl_adreno_xmem_attn_state adreno_xmem_attn;
#endif
+ cl_kernel kernel_top_k_f32, kernel_top_k_merge_f32;
+ cl_kernel kernel_lightning_indexer;
cl_kernel kernel_get_rows_f32, kernel_get_rows_f16, kernel_get_rows_i32, 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;
@@ -2091,6 +2095,37 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
GGML_LOG_CONT(".");
}
+ // lightning_indexer
+ {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "lightning_indexer.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("lightning_indexer.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_lightning_indexer = clCreateKernel(prog, "kernel_lightning_indexer", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
+ // top_k
+ {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "top_k.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("top_k.cl");
+#endif
+ cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+ CL_CHECK((backend_ctx->kernel_top_k_f32 = clCreateKernel(prog, "kernel_top_k_f32", &err), err));
+ CL_CHECK((backend_ctx->kernel_top_k_merge_f32 = clCreateKernel(prog, "kernel_top_k_merge_f32", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ GGML_LOG_CONT(".");
+ }
+
// get_rows
{
#ifdef GGML_OPENCL_EMBED_KERNELS
@@ -7437,9 +7472,8 @@ static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev) {
// subgroup shuffle support (N_SPLIT>1 FA kernel)
backend_ctx->has_qcom_subgroup_shuffle = strstr(ext_buffer.data(), "cl_qcom_subgroup_shuffle") != NULL;
- backend_ctx->has_subgroup_shuffle =
- strstr(ext_buffer.data(), "cl_khr_subgroup_shuffle") != NULL ||
- backend_ctx->has_qcom_subgroup_shuffle;
+ backend_ctx->has_khr_subgroup_shuffle = strstr(ext_buffer.data(), "cl_khr_subgroup_shuffle") != NULL;
+ backend_ctx->has_subgroup_shuffle = backend_ctx->has_khr_subgroup_shuffle || backend_ctx->has_qcom_subgroup_shuffle;
// check for cl_khr_integer_dot_product
// cl_qcom_dot_product8 uses signed * unsigned
@@ -9943,6 +9977,14 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_I32 &&
ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op) &&
backend_ctx->get_kernel_workgroup_size(backend_ctx->kernel_argmax_f32_i32) >= 64;
+ case GGML_OP_LIGHTNING_INDEXER:
+ return op->src[0]->type == GGML_TYPE_F32 &&
+ (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_BF16) &&
+ op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F16 &&
+ op->src[0]->nb[0] == sizeof(float) && op->src[1]->nb[0] == ggml_type_size(op->src[1]->type) &&
+ op->src[2]->nb[0] == sizeof(float) && op->src[3]->nb[0] == sizeof(ggml_fp16_t);
+ case GGML_OP_TOP_K:
+ return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_I32;
case GGML_OP_ARGSORT: {
load_cl_kernels_argsort(backend_ctx);
@@ -30991,6 +31033,157 @@ static void ggml_cl_argmax(ggml_backend_t backend, const ggml_tensor * src0, con
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
}
+static void ggml_cl_lightning_indexer(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+ ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+ cl_kernel kernel = backend_ctx->kernel_lightning_indexer;
+ const ggml_tensor * src2 = dst->src[2];
+ const ggml_tensor * src3 = dst->src[3];
+ ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *) src0->extra;
+ ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *) src1->extra;
+ ggml_tensor_extra_cl * extra2 = (ggml_tensor_extra_cl *) src2->extra;
+ ggml_tensor_extra_cl * extra3 = (ggml_tensor_extra_cl *) src3->extra;
+ ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *) dst->extra;
+
+ const cl_ulong offset0 = extra0->offset + src0->view_offs;
+ const cl_ulong offset1 = extra1->offset + src1->view_offs;
+ const cl_ulong offset2 = extra2->offset + src2->view_offs;
+ const cl_ulong offset3 = extra3->offset + src3->view_offs;
+ const cl_ulong offsetd = extrad->offset + dst->view_offs;
+
+ const int head_size = src0->ne[0];
+ const int n_heads = src0->ne[1];
+ const int n_tokens = src0->ne[2];
+ const int n_keys = src1->ne[2];
+ const int n_masks = src3->ne[3];
+ // Keep the local-memory reduction for small grids or few heads.
+ const int use_shuffle = head_size == 128 && backend_ctx->has_khr_subgroup_shuffle && n_heads >= 32 && ggml_nelements(dst) >= 1024;
+ // BF16 tensors are stored as F16 on the device.
+ const int type_k = src1->type == GGML_TYPE_F32 ? 0 : 1;
+
+ const cl_ulong nbq1 = src0->nb[1];
+ const cl_ulong nbq2 = src0->nb[2];
+ const cl_ulong nbq3 = src0->nb[3];
+ const cl_ulong nbk2 = src1->nb[2];
+ const cl_ulong nbk3 = src1->nb[3];
+ const cl_ulong nbw1 = src2->nb[1];
+ const cl_ulong nbw3 = src2->nb[3];
+ const cl_ulong nbm1 = src3->nb[1];
+ const cl_ulong nbm3 = src3->nb[3];
+
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra0->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extra1->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset1));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &extra2->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offset2));
+ CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_mem), &extra3->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &offset3));
+ CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, 10, sizeof(int), &head_size));
+ CL_CHECK(clSetKernelArg(kernel, 11, sizeof(int), &n_heads));
+ CL_CHECK(clSetKernelArg(kernel, 12, sizeof(int), &n_tokens));
+ CL_CHECK(clSetKernelArg(kernel, 13, sizeof(int), &n_keys));
+ CL_CHECK(clSetKernelArg(kernel, 14, sizeof(int), &n_masks));
+ CL_CHECK(clSetKernelArg(kernel, 15, sizeof(int), &type_k));
+ CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &nbq1));
+ CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &nbq2));
+ CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &nbq3));
+ CL_CHECK(clSetKernelArg(kernel, 19, sizeof(cl_ulong), &nbk2));
+ CL_CHECK(clSetKernelArg(kernel, 20, sizeof(cl_ulong), &nbk3));
+ CL_CHECK(clSetKernelArg(kernel, 21, sizeof(cl_ulong), &nbw1));
+ CL_CHECK(clSetKernelArg(kernel, 22, sizeof(cl_ulong), &nbw3));
+ CL_CHECK(clSetKernelArg(kernel, 23, sizeof(cl_ulong), &nbm1));
+ CL_CHECK(clSetKernelArg(kernel, 24, sizeof(cl_ulong), &nbm3));
+ CL_CHECK(clSetKernelArg(kernel, 25, sizeof(int), &use_shuffle));
+
+ const int keys_per_group = use_shuffle ? 16 : 4;
+ size_t global_work_size[] = {(size_t) ((src1->ne[2] + keys_per_group - 1)/keys_per_group)*64, (size_t) src0->ne[2], (size_t) src0->ne[3]};
+ size_t local_work_size[] = {64, 1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+}
+
+static void ggml_cl_top_k(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+ GGML_UNUSED(src1);
+
+ ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+
+ ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *) src0->extra;
+ ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *) dst->extra;
+
+ const cl_ulong offset0 = extra0->offset + src0->view_offs;
+ const cl_ulong offsetd = extrad->offset + dst->view_offs;
+
+ GGML_TENSOR_LOCALS(int, ne0, src0, ne);
+ GGML_TENSOR_LOCALS(cl_ulong, nb0, src0, nb);
+
+ const int k = dst->ne[0];
+ cl_kernel kernel = backend_ctx->kernel_top_k_f32;
+ const int max_wg = backend_ctx->get_kernel_workgroup_size(kernel);
+
+ int nth = 1;
+ const int target = MIN(512, MAX(32, ne00/4));
+ while (2*nth <= MIN(target, max_wg)) {
+ nth *= 2;
+ }
+
+ const size_t n_rows = ggml_nrows(src0);
+ const bool tiled = ne00 > 32768 && k > 0 && k <= 256 && n_rows <= 4;
+ const int tile_width = tiled ? 4096 : ne00;
+ const int n_tiles = tiled ? 1 + (ne00 - 1)/tile_width : 1;
+ const int n_candidates = (n_tiles - 1)*k + MIN(k, ne00 - (n_tiles - 1)*tile_width);
+ cl_mem candidates = nullptr;
+ if (tiled) {
+ GGML_ASSERT(n_rows <= SIZE_MAX/sizeof(cl_uint2)/(size_t) n_candidates);
+ backend_ctx->prealloc_top_k.allocate(backend_ctx->context, n_rows*(size_t) n_candidates*sizeof(cl_uint2));
+ candidates = backend_ctx->prealloc_top_k.buffer;
+ }
+
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra0->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &ne02));
+ CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &nb00));
+ CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_ulong), &nb01));
+ CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &nb02));
+ CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_ulong), &nb03));
+ CL_CHECK(clSetKernelArg(kernel, 11, sizeof(int), &k));
+ CL_CHECK(clSetKernelArg(kernel, 12, MAX(259, 2*nth)*sizeof(cl_uint), nullptr));
+ CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_mem), &candidates));
+ CL_CHECK(clSetKernelArg(kernel, 14, sizeof(int), &tile_width));
+
+ size_t global_work_size[] = {(size_t) nth*n_tiles, n_rows, 1};
+ size_t local_work_size[] = {(size_t) nth, 1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size, local_work_size, dst);
+
+ if (tiled) {
+ // The global top-k is contained in the top-k of each tile.
+ kernel = backend_ctx->kernel_top_k_merge_f32;
+ const int max_wg_merge = backend_ctx->get_kernel_workgroup_size(kernel);
+
+ int nth_merge = 1;
+ const int target_merge = MIN(512, MAX(32, n_candidates/4));
+ while (2*nth_merge <= MIN(target_merge, max_wg_merge)) {
+ nth_merge *= 2;
+ }
+
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &candidates));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(int), &n_candidates));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &k));
+ CL_CHECK(clSetKernelArg(kernel, 5, MAX(259, 2*nth_merge)*sizeof(cl_uint), nullptr));
+
+ global_work_size[0] = nth_merge;
+ local_work_size[0] = nth_merge;
+
+ backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size, local_work_size, dst);
+ }
+}
+
static void ggml_cl_argsort(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
GGML_ASSERT(src0);
GGML_ASSERT(src0->extra);
@@ -31910,6 +32103,12 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
}
func = ggml_cl_argmax;
break;
+ case GGML_OP_LIGHTNING_INDEXER:
+ func = ggml_cl_lightning_indexer;
+ break;
+ case GGML_OP_TOP_K:
+ func = ggml_cl_top_k;
+ break;
case GGML_OP_ARGSORT:
if (!any_on_device) {
return false;
diff --git a/ggml/src/ggml-opencl/kernels/lightning_indexer.cl b/ggml/src/ggml-opencl/kernels/lightning_indexer.cl
new file mode 100644
index 000000000..efd2f1165
--- /dev/null
+++ b/ggml/src/ggml-opencl/kernels/lightning_indexer.cl
@@ -0,0 +1,111 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#ifdef cl_khr_subgroup_shuffle
+#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
+#endif
+
+kernel void kernel_lightning_indexer(
+ global const char * q,
+ ulong offset_q,
+ global const char * k,
+ ulong offset_k,
+ global const char * w,
+ ulong offset_w,
+ global const char * mask,
+ ulong offset_m,
+ global float * dst,
+ ulong offset_d,
+ int head_size,
+ int n_heads,
+ int n_tokens,
+ int n_keys,
+ int n_masks,
+ int type_k,
+ ulong nbq1,
+ ulong nbq2,
+ ulong nbq3,
+ ulong nbk2,
+ ulong nbk3,
+ ulong nbw1,
+ ulong nbw3,
+ ulong nbm1,
+ ulong nbm3,
+ int use_shuffle
+) {
+ const int tid = get_local_id(0);
+ const int lane = tid & (use_shuffle ? 3 : 15);
+ const int key = use_shuffle ? get_group_id(0)*16 + tid/4 : get_group_id(0)*4 + tid/16;
+ const int token = get_group_id(1);
+ const int stream = get_group_id(2);
+
+ q += offset_q + token*nbq2 + stream*nbq3;
+ k += offset_k + (ulong) key*nbk2 + stream*nbk3;
+ w += offset_w + token*nbw1 + stream*nbw3;
+ mask += offset_m + token*nbm1 + (stream % n_masks)*nbm3;
+ dst = (global float *) ((global char *) dst + offset_d);
+
+#ifdef cl_khr_subgroup_shuffle
+ if (use_shuffle) {
+ // Each group of four lanes keeps one key for all query heads.
+ float4 cached[8];
+ #pragma unroll
+ for (int j = 0; j < 8; ++j) {
+ const int i = 4*(lane + j*4);
+ cached[j] = key < n_keys ? (type_k == 0 ? vload4(0, (global const float *) k + i) : convert_float4(vload4(0, (global const half *) k + i))) : (float4) (0.0f);
+ }
+
+ float score = 0.0f;
+ for (int h = 0; h < n_heads; ++h) {
+ global const float * q_head = (global const float *) (q + h*nbq1);
+ float4 sum = (float4) (0.0f);
+ #pragma unroll
+ for (int j = 0; j < 8; ++j) {
+ sum = mad(vload4(0, q_head + 4*(lane + j*4)), cached[j], sum);
+ }
+ float value = sum.s0 + sum.s1 + sum.s2 + sum.s3;
+ value += sub_group_shuffle_xor(value, 2);
+ value += sub_group_shuffle_xor(value, 1);
+ if (lane == 0) {
+ score = mad(fmax(value, 0.0f), ((global const float *) w)[h], score);
+ }
+ }
+ if (lane == 0 && key < n_keys) {
+ dst[((ulong) stream*n_tokens + token)*n_keys + key] = score + (float) ((global const half *) mask)[key];
+ }
+ return;
+ }
+#endif
+
+ local float sums[64];
+ float score = 0.0f;
+ for (int h = 0; h < n_heads; ++h) {
+ global const float * q_head = (global const float *) (q + h*nbq1);
+ float sum = 0.0f;
+ if (key < n_keys) {
+ for (int i = lane; i < head_size; i += 16) {
+ float kv;
+ if (type_k == 0) {
+ kv = ((global const float *) k)[i];
+ } else {
+ kv = (float) ((global const half *) k)[i];
+ }
+ sum = mad(q_head[i], kv, sum);
+ }
+ }
+ sums[tid] = sum;
+ barrier(CLK_LOCAL_MEM_FENCE);
+ for (int stride = 8; stride > 0; stride /= 2) {
+ if (lane < stride) {
+ sums[tid] += sums[tid + stride];
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+ if (lane == 0) {
+ score = mad(fmax(sums[tid], 0.0f), ((global const float *) w)[h], score);
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+
+ if (lane == 0 && key < n_keys) {
+ dst[((ulong) stream*n_tokens + token)*n_keys + key] = score + (float) ((global const half *) mask)[key];
+ }
+}
diff --git a/ggml/src/ggml-opencl/kernels/top_k.cl b/ggml/src/ggml-opencl/kernels/top_k.cl
new file mode 100644
index 000000000..6b5126e0d
--- /dev/null
+++ b/ggml/src/ggml-opencl/kernels/top_k.cl
@@ -0,0 +1,207 @@
+#ifdef cl_khr_subgroups
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#endif
+
+uint top_k_key(float value) {
+ uint bits = as_uint(value);
+ if ((bits & 0x7fffffffu) == 0) {
+ bits = 0;
+ }
+ return bits & 0x80000000u ? ~bits : bits ^ 0x80000000u;
+}
+
+uint4 top_k_keys(global const char * src, int i, int n, ulong nb0) {
+ if (nb0 == sizeof(float) && i + 3 < n) {
+ const float4 v = vload4(0, (global const float *) src + i);
+ return (uint4) (top_k_key(v.s0), top_k_key(v.s1), top_k_key(v.s2), top_k_key(v.s3));
+ }
+ return (uint4) (
+ i < n ? top_k_key(*(global const float *) (src + (ulong) i *nb0)) : 0,
+ i + 1 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 1)*nb0)) : 0,
+ i + 2 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 2)*nb0)) : 0,
+ i + 3 < n ? top_k_key(*(global const float *) (src + (ulong) (i + 3)*nb0)) : 0);
+}
+
+void top_k_store(global int * dst, global uint2 * candidates, global const char * src,
+ ulong nb0, int indexed, int base, int i, uint pos) {
+ const int index = indexed ? *(global const int *) (src + (ulong) i*nb0 + sizeof(float)) : base + i;
+ if (candidates) {
+ candidates[pos] = (uint2) (as_uint(*(global const float *) (src + (ulong) i*nb0)), index);
+ } else {
+ dst[pos] = index;
+ }
+}
+
+void top_k_select(
+ global const char * src,
+ global int * dst,
+ global uint2 * candidates,
+ ulong nb00,
+ int ne00,
+ int k,
+ int base,
+ int indexed,
+ local uint * hist
+) {
+ const int tid = get_local_id(0);
+ const int nth = get_local_size(0);
+ if (k == ne00) {
+ for (int i = tid; i < ne00; i += nth) {
+ top_k_store(dst, candidates, src, nb00, indexed, base, i, i);
+ }
+ return;
+ }
+
+ uint prefix = 0;
+ uint mask = 0;
+ uint rank = k;
+ for (int shift = 24; shift >= 0; shift -= 8) {
+ for (int i = tid; i < 256; i += nth) {
+ hist[i] = 0;
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ for (int i = 4*tid; i < ne00; i += 4*nth) {
+ const uint4 v = top_k_keys(src, i, ne00, nb00);
+ const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+ #pragma unroll
+ for (int j = 0; j < 4; ++j) {
+ const uint key = keys[j];
+ if (i + j < ne00 && (key & mask) == prefix) {
+ atomic_inc(hist + ((key >> shift) & 255));
+ }
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ if (tid == 0) {
+ uint above = 0;
+ for (int b = 255; b >= 0; --b) {
+ uint count = hist[b];
+ if (above + count >= rank) {
+ hist[256] = b;
+ hist[257] = rank - above;
+ hist[258] = count;
+ break;
+ }
+ above += count;
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ prefix |= hist[256] << shift;
+ rank = hist[257];
+ const bool done = hist[258] == rank;
+ barrier(CLK_LOCAL_MEM_FENCE);
+ if (done) {
+ // The whole bucket is selected, so its low bits do not matter.
+ break;
+ }
+ mask |= 255u << shift;
+ }
+
+ uint greater = 0;
+ uint equal = 0;
+ for (int i = 4*tid; i < ne00; i += 4*nth) {
+ const uint4 v = top_k_keys(src, i, ne00, nb00);
+ const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+ #pragma unroll
+ for (int j = 0; j < 4; ++j) {
+ greater += i + j < ne00 && keys[j] > prefix;
+ equal += i + j < ne00 && keys[j] == prefix;
+ }
+ }
+ uint pos_greater;
+ uint pos_equal;
+#ifdef cl_khr_subgroups
+ const uint sg = get_sub_group_id();
+ const uint nsg = get_num_sub_groups();
+ const uint total_greater = sub_group_reduce_add(greater);
+ const uint total_equal = sub_group_reduce_add(equal);
+ pos_greater = sub_group_scan_exclusive_add(greater);
+ pos_equal = sub_group_scan_exclusive_add(equal);
+ if (get_sub_group_local_id() == 0) {
+ hist[sg] = total_greater;
+ hist[nsg + sg] = total_equal;
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ for (uint i = 0; i < nsg; ++i) {
+ pos_equal += hist[i];
+ if (i < sg) {
+ pos_greater += hist[i];
+ pos_equal += hist[nsg + i];
+ }
+ }
+#else
+ hist[tid] = greater;
+ hist[nth + tid] = equal;
+ barrier(CLK_LOCAL_MEM_FENCE);
+ pos_greater = 0;
+ pos_equal = 0;
+ for (int i = 0; i < nth; ++i) {
+ pos_equal += hist[i];
+ if (i < tid) {
+ pos_greater += hist[i];
+ pos_equal += hist[nth + i];
+ }
+ }
+#endif
+ for (int i = 4*tid; i < ne00; i += 4*nth) {
+ const uint4 v = top_k_keys(src, i, ne00, nb00);
+ const uint keys[4] = {v.s0, v.s1, v.s2, v.s3};
+ #pragma unroll
+ for (int j = 0; j < 4; ++j) {
+ if (i + j < ne00) {
+ if (keys[j] > prefix) {
+ top_k_store(dst, candidates, src, nb00, indexed, base, i + j, pos_greater++);
+ } else if (keys[j] == prefix && pos_equal < k) {
+ top_k_store(dst, candidates, src, nb00, indexed, base, i + j, pos_equal++);
+ }
+ }
+ }
+ }
+}
+
+kernel void kernel_top_k_f32(
+ global const char * src,
+ ulong offset0,
+ global int * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int k_max,
+ local uint * hist,
+ global uint2 * candidates,
+ int tile_width
+) {
+ const ulong row = get_group_id(1);
+ const int part = get_group_id(0);
+ const int parts = get_num_groups(0);
+ const int base = part*tile_width;
+ const int n = min(ne00 - base, tile_width);
+ const int k = min(k_max, n);
+ const int out_width = (parts - 1)*k_max + min(k_max, ne00 - (parts - 1)*tile_width);
+ src += offset0 + (row % ne01)*nb01 + (row / ne01 % ne02)*nb02 + (row / ne01 / ne02)*nb03 + (ulong) base*nb00;
+ if (candidates) {
+ candidates += row*out_width + part*k_max;
+ } else {
+ dst = (global int *) ((global char *) dst + offsetd) + row*out_width + part*k_max;
+ }
+ top_k_select(src, dst, candidates, nb00, n, k, base, 0, hist);
+}
+
+kernel void kernel_top_k_merge_f32(
+ global const uint2 * candidates,
+ global int * dst,
+ ulong offsetd,
+ int n_candidates,
+ int k,
+ local uint * hist
+) {
+ const ulong row = get_group_id(1);
+ candidates += row*n_candidates;
+ dst = (global int *) ((global char *) dst + offsetd) + row*k;
+ top_k_select((global const char *) candidates, dst, 0, sizeof(uint2), n_candidates, k, 0, 1, hist);
+}