public_git
/
whisper.cpp
mirror of https://github.com/ggml-org/whisper.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

kleidiai : add SME2 F32 GEMV kernel support (llama/26891)

This commit is contained in:
Charles Xu 2026-08-21 10:33:30 +02:00 committed by Georgi Gerganov
parent d60ef65077
commit d6c416e222
3 changed files with 70 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
ggml/src/ggml-cpu/CMakeLists.txt
View File

@ -639,6 +639,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/)
set(ARCH_FLAGS_TEMP "${ARCH_FLAGS}")
@ -701,6 +702,8 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa_asm.S
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa_asm.S
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla_asm.S
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_bf16p2vlx2_f32_sme.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.c
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f16pmrx2_f32_neon.c

48
ggml/src/ggml-cpu/kleidiai/kernels.cpp
View File

@ -23,6 +23,7 @@
#include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h"
#include "kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h"
#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
#include "kai_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla.h"
#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa.h"
#include "kai_lhs_pack_bf16p2vlx2_f32_sme.h"
@ -76,6 +77,21 @@ static inline void kernel_run_fn10(size_t m, size_t n, size_t k, size_t /*bl*/,
Fn(m, n, k, lhs, rhs, dst, dst_stride_row, dst_stride_col, clamp_min, clamp_max);
}
template <void (*Fn)(size_t, size_t, size_t, const void *, size_t, const void *, void *, size_t, size_t, float, float)>
static inline void kernel_run_lhs_stride_fn10(size_t m,
size_t n,
size_t k,
size_t lhs_stride,
const void * lhs,
const void * rhs,
void * dst,
size_t dst_stride_row,
size_t dst_stride_col,
float clamp_min,
float clamp_max) {
Fn(m, n, k, lhs, lhs_stride, rhs, dst, dst_stride_row, dst_stride_col, clamp_min, clamp_max);
}
template<void(*Fn)(size_t,size_t,size_t,const void*,const void*,float*,size_t,size_t,float,float)>
static inline void kernel_run_float_fn10(size_t m, size_t n, size_t k, size_t /*bl*/,
const void* lhs, const void* rhs, void* dst,
@ -947,25 +963,25 @@ static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
/* .packed_size_ex = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
/* .pack_func_ex = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
},
/* SME GEMV */
/* SME2 GEMV */
{
/* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
/* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
/* .get_mr = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
/* .get_nr = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
/* .get_kr = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
/* .get_sr = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
/* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
/* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
/* .get_lhs_offset_ex = */ nullptr,
/* .get_rhs_packed_offset_ex = */ nullptr,
/* .run_kernel_ex = */ nullptr,
/* .get_m_step = */ kai_get_m_step_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla,
/* .get_n_step = */ kai_get_n_step_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla,
/* .get_mr = */ kai_get_m_step_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla,
/* .get_nr = */ kai_get_nr_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla,
/* .get_kr = */ kai_get_kr_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla,
/* .get_sr = */ kai_get_sr_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla,
/* .get_dst_offset = */ kai_get_dst_offset_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla,
/* .get_dst_size = */ kai_get_dst_size_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla,
/* .get_lhs_offset_ex = */ &kernel_offs_fn2<kai_get_lhs_offset_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla>,
/* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla>,
/* .run_kernel_ex = */ &kernel_run_lhs_stride_fn10<kai_run_matmul_clamp_f32_f32_f32p2vlx1b_1x16vl_sme2_mla>,
},
/* .gemv_lhs_info = */ {
/* .get_offset = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
/* .get_packed_offset_ex = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
/* .packed_size_ex = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
/* .pack_func_ex = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
/* .get_offset = */ nullptr,
/* .get_packed_offset_ex = */ nullptr,
/* .packed_size_ex = */ nullptr,
/* .pack_func_ex = */ nullptr,
},
/* .rhs_info = */ {
/* .packed_stride = */ nullptr,

48
ggml/src/ggml-cpu/kleidiai/kleidiai.cpp
View File

@ -696,6 +696,15 @@ class tensor_traits : public ggml::cpu::tensor_traits {
}
if (op->src[0]->type == GGML_TYPE_F32) {
ggml_kleidiai_kernels * primary = kernel_chain[0];
kernel_info * gemv_kernel = primary ? &primary->gemv : nullptr;
if (is_gemv && op->src[1]->nb[0] == (int64_t) sizeof(float) && gemv_kernel &&
gemv_kernel->get_lhs_offset_ex && gemv_kernel->get_rhs_packed_offset_ex &&
gemv_kernel->run_kernel_ex && gemv_kernel->get_dst_offset) {
size = 0;
return true;
}
size_t cursor = 0;
bool any_slot = false;
@ -811,15 +820,28 @@ class tensor_traits : public ggml::cpu::tensor_traits {
return false;
}
kernel_info * kernel = &kernels->gemm;
const size_t k = ne00;
const size_t m = ne11;
const size_t n = ne01;
const bool use_gemv = m == 1 && src1->nb[0] == (int64_t) sizeof(float) &&
kernels->gemv.get_lhs_offset_ex &&
kernels->gemv.get_rhs_packed_offset_ex &&
kernels->gemv.run_kernel_ex &&
kernels->gemv.get_dst_offset;
kernel_info * kernel = use_gemv ? &kernels->gemv : &kernels->gemm;
lhs_packing_info * lhs_info = &kernels->gemm_lhs_info;
if (!kernel || !lhs_info || !lhs_info->get_offset || !lhs_info->get_packed_offset_ex ||
!lhs_info->packed_size_ex || !lhs_info->pack_func_ex ||
if (!kernel || !kernel->get_lhs_offset_ex ||
!kernel->get_rhs_packed_offset_ex || !kernel->run_kernel_ex || !kernel->get_dst_offset) {
return false;
}
if (!use_gemv && (!lhs_info || !lhs_info->get_offset || !lhs_info->get_packed_offset_ex ||
!lhs_info->packed_size_ex || !lhs_info->pack_func_ex)) {
return false;
}
const kleidiai_weight_header * header = kleidiai_weight_header_from_ptr(src0->data);
const bool has_header = kleidiai_is_weight_header_valid(header);
@ -832,16 +854,14 @@ class tensor_traits : public ggml::cpu::tensor_traits {
const int nth = params->nth > 0 ? params->nth : 1;
const int ith = params->ith;
const size_t k = ne00;
const size_t m = ne11;
const size_t n = ne01;
const size_t mr = kernel->get_mr();
const size_t kr = kernel->get_kr();
const size_t sr = kernel->get_sr();
const size_t lhs_packed_size = lhs_info->packed_size_ex(m, k, 0, mr, kr, sr);
GGML_ASSERT(lhs_packed_size <= params->wsize);
const size_t lhs_packed_size = use_gemv ? 0 : lhs_info->packed_size_ex(m, k, 0, mr, kr, sr);
if (!use_gemv) {
GGML_ASSERT(lhs_packed_size <= params->wsize);
}
uint8_t * lhs_packed = static_cast<uint8_t *>(params->wdata);
const size_t dst_stride = dst->nb[1];
@ -853,7 +873,7 @@ class tensor_traits : public ggml::cpu::tensor_traits {
const uint8_t * lhs_batch_base = static_cast<const uint8_t *>(src1->data) + batch_idx * src1->nb[2];
uint8_t * dst_batch_base = static_cast<uint8_t *>(dst->data) + batch_idx * dst->nb[2];
{
if (!use_gemv) {
const int64_t m_roundup_mr = kai_roundup((int64_t)m, (int64_t)mr);
int64_t max_threads = mr ? (m_roundup_mr / (int64_t)mr) : nth;
max_threads = std::max<int64_t>(1, max_threads);
@ -903,15 +923,17 @@ class tensor_traits : public ggml::cpu::tensor_traits {
const size_t n_to_process = std::min(chunk_cols, n - n_start);
if (n_to_process > 0) {
const size_t lhs_packed_offset = lhs_info->get_packed_offset_ex(0, k, 0, mr, kr, sr);
const size_t lhs_offset = use_gemv ? kernel->get_lhs_offset_ex(0, k, 0)
: lhs_info->get_packed_offset_ex(0, k, 0, mr, kr, sr);
const size_t rhs_packed_offset = kernel->get_rhs_packed_offset_ex(n_start, k, 0);
const size_t dst_offset = kernel->get_dst_offset(0, n_start, dst_stride);
const void * lhs_ptr = lhs_packed + lhs_packed_offset;
const void * lhs_ptr = use_gemv ? lhs_batch_base + lhs_offset
: lhs_packed + lhs_offset;
const void * rhs_ptr = rhs_base + rhs_packed_offset;
float * dst_ptr = reinterpret_cast<float *>(dst_batch_base + dst_offset);
kernel->run_kernel_ex(m, n_to_process, k, 0,
kernel->run_kernel_ex(m, n_to_process, k, use_gemv ? src1->nb[1] : 0,
lhs_ptr,
rhs_ptr,
dst_ptr,