opencl: add F32-to-I32 ARGMAX (#30897)

* opencl: add F32-to-I32 ARGMAX

Assisted-by: OpenCode

* opencl: init to 0 to be consistent with cpu

---------

Co-authored-by: Li He <lih@qti.qualcomm.com>

ggml/src/ggml-opencl/CMakeLists.txt

@@ -64,6 +64,7 @@ set(GGML_OPENCL_KERNELS
     add
     add_id
     moe_add_id_glu
+    argmax
     argsort
     tri
     fill

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

@@ -646,8 +646,8 @@ struct ggml_backend_opencl_context {
     ggml_cl_version platform_version;
     ggml_cl_version opencl_c_version;
 
-    // argsort is loaded in supports_op because its availability depends on how
-    // many workgroups are allowed, which requires kernel compilation.
+    // argmax and argsort are loaded in supports_op; argsort availability depends on the compiled kernel's workgroup limit.
+    bool kernels_loaded_argmax = false;
     bool kernels_loaded_argsort = false;
     // rest of the kernels are currently always loaded in alloc_buffer.
     bool kernels_loaded = false;
@@ -972,6 +972,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_solve_tri_f32;
     cl_kernel kernel_im2col_f32, kernel_im2col_f16;
     cl_kernel kernel_im2col_3d_f32, kernel_im2col_3d_f16;
+    cl_kernel kernel_argmax_f32_i32;
     cl_kernel kernel_argsort_f32_i32;
     cl_kernel kernel_sum_rows_f32, kernel_sum_rows_f32_4;
     cl_kernel kernel_cumsum_blk, kernel_cumsum_add;
@@ -1485,6 +1486,27 @@ static cl_program build_program_from_binary(cl_context ctx, cl_device_id dev, co
     return p;
 }
 
+static void load_cl_kernels_argmax(ggml_backend_opencl_context *backend_ctx) {
+    auto opencl_c_std =
+        std::string("CL") + std::to_string(backend_ctx->opencl_c_version.major) + "." + std::to_string(backend_ctx->opencl_c_version.minor);
+    std::string compile_opts = std::string("-cl-std=") + opencl_c_std;
+
+    if (!backend_ctx->kernels_loaded_argmax) {
+        cl_int err;
+#ifdef GGML_OPENCL_EMBED_KERNELS
+        const std::string kernel_src {
+            #include "argmax.cl.h"
+        };
+#else
+        const std::string kernel_src = read_file("argmax.cl");
+#endif
+        cl_program prog = build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
+        CL_CHECK((backend_ctx->kernel_argmax_f32_i32 = clCreateKernel(prog, "kernel_argmax_f32_i32", &err), err));
+        CL_CHECK(clReleaseProgram(prog));
+        backend_ctx->kernels_loaded_argmax = true;
+    }
+}
+
 static void load_cl_kernels_argsort(ggml_backend_opencl_context *backend_ctx) {
     // compiler options for general kernels
     auto opencl_c_std =
@@ -9693,6 +9715,11 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
             return true;
         case GGML_OP_IM2COL_3D:
             return true;
+        case GGML_OP_ARGMAX:
+            load_cl_kernels_argmax(backend_ctx);
+            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_ARGSORT: {
             load_cl_kernels_argsort(backend_ctx);
 
@@ -30164,6 +30191,41 @@ static void ggml_cl_im2col_3d(ggml_backend_t backend, const ggml_tensor * src0,
     backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
 }
 
+static void ggml_cl_argmax(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_ASSERT(src0);
+    GGML_ASSERT(src0->extra);
+    GGML_ASSERT(dst);
+    GGML_ASSERT(dst->extra);
+    GGML_UNUSED(src1);
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_I32);
+    GGML_ASSERT(ggml_is_contiguous(src0));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    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;
+    const int ne00 = src0->ne[0];
+    const int nrows = ggml_nrows(src0);
+    constexpr int nth = 64;
+
+    cl_kernel kernel = backend_ctx->kernel_argmax_f32_i32;
+    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),      &nrows));
+
+    size_t global_work_size[] = { (size_t) nth, (size_t) CEIL_DIV(nrows, 2), 1 };
+    size_t local_work_size[] = { (size_t) nth, 1, 1 };
+    backend_ctx->enqueue_ndrange_kernel(kernel, 3, 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);
@@ -31071,6 +31133,12 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
             }
             func = ggml_cl_im2col_3d;
             break;
+        case GGML_OP_ARGMAX:
+            if (!any_on_device) {
+                return false;
+            }
+            func = ggml_cl_argmax;
+            break;
         case GGML_OP_ARGSORT:
             if (!any_on_device) {
                 return false;

ggml/src/ggml-opencl/kernels/argmax.cl

@@ -0,0 +1,55 @@
+#define ARGMAX_LANES_PER_ROW 32
+#define ARGMAX_ROWS_PER_GROUP 2
+#define ARGMAX_BLOCK_SIZE (ARGMAX_LANES_PER_ROW * ARGMAX_ROWS_PER_GROUP)
+
+kernel void kernel_argmax_f32_i32(
+        global const char * src0,
+        ulong               offset0,
+        global char *       dst,
+        ulong               offsetd,
+        int                 ne00,
+        int                 nrows) {
+    const int local_id = get_local_id(0);
+    const int col = local_id % ARGMAX_LANES_PER_ROW;
+    const int row = get_group_id(1) * ARGMAX_ROWS_PER_GROUP + local_id / ARGMAX_LANES_PER_ROW;
+
+    float max_value = -INFINITY;
+    int max_index = 0;
+
+    if (row < nrows) {
+        global const float * src_row = (global const float *) (src0 + offset0) + row * ne00;
+        for (int i = col; i < ne00; i += ARGMAX_LANES_PER_ROW) {
+            const float value = src_row[i];
+            if (value > max_value || (value == max_value && i > max_index)) {
+                max_value = value;
+                max_index = i;
+            }
+        }
+    }
+
+    local float max_values[ARGMAX_BLOCK_SIZE];
+    local int max_indices[ARGMAX_BLOCK_SIZE];
+
+    max_values[local_id] = max_value;
+    max_indices[local_id] = max_index;
+
+    barrier(CLK_LOCAL_MEM_FENCE);
+
+    const int row_offset = local_id / ARGMAX_LANES_PER_ROW * ARGMAX_LANES_PER_ROW;
+    for (int stride = ARGMAX_LANES_PER_ROW / 2; stride > 0; stride /= 2) {
+        if (col < stride) {
+            const float value = max_values[row_offset + col + stride];
+            const int index = max_indices[row_offset + col + stride];
+            if (value > max_values[row_offset + col] || (value == max_values[row_offset + col] && index > max_indices[row_offset + col])) {
+                max_values[row_offset + col] = value;
+                max_indices[row_offset + col] = index;
+            }
+        }
+        barrier(CLK_LOCAL_MEM_FENCE);
+    }
+
+    if (col == 0 && row < nrows) {
+        global int * dst_row = (global int *) (dst + offsetd) + row;
+        dst_row[0] = max_indices[row_offset];
+    }
+}