diff --git a/ggml/src/ggml-opencl/CMakeLists.txt b/ggml/src/ggml-opencl/CMakeLists.txt index 917578c9f..977013b94 100644 --- a/ggml/src/ggml-opencl/CMakeLists.txt +++ b/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 diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 2312c6c7c..d56b3488d 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/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; diff --git a/ggml/src/ggml-opencl/kernels/argmax.cl b/ggml/src/ggml-opencl/kernels/argmax.cl new file mode 100644 index 000000000..24ed0e68b --- /dev/null +++ b/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]; + } +}