Skip to content

Commit c5d91a7

Browse files
chaxu01chaxu01
authored
ggml-cpu: Add CPU backend support for KleidiAI library (ggml-org#11390)
* ggml-cpu: Add CPU backend support for KleidiAI library * Add environmental variable GGML_KLEIDIAI_SME * Add support for multithread LHS conversion * Switch kernel selection order to dotprod and i8mm * updates for review comments * More updates for review comments * Reorganize and rename KleidiAI files * Move ggml-cpu-traits.h to source file * Update cmake for SME build and add alignment for SME * Remove append GGML_USE_CPU_KLEIDIAI to the GGML_CDEF_PUBLIC list
1 parent 4806498 commit c5d91a7

File tree

9 files changed

+767
-10
lines changed

9 files changed

+767
-10
lines changed

‎ggml/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ endif()
102102

103103
option(GGML_CPU_HBM "ggml: use memkind for CPU HBM" OFF)
104104
option(GGML_CPU_AARCH64 "ggml: use runtime weight conversion of Q4_0 to Q4_X_X" ON)
105+
option(GGML_CPU_KLEIDIAI "ggml: use KleidiAI optimized kernels if applicable" OFF)
105106
option(GGML_AVX "ggml: enable AVX" ${INS_ENB})
106107
option(GGML_AVX_VNNI "ggml: enable AVX-VNNI" OFF)
107108
option(GGML_AVX2 "ggml: enable AVX2" ${INS_ENB})

‎ggml/include/ggml-cpu.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ extern "C" {
9595
GGML_BACKEND_API int ggml_cpu_has_matmul_int8(void);
9696
GGML_BACKEND_API int ggml_cpu_has_sve (void);
9797
GGML_BACKEND_API int ggml_cpu_get_sve_cnt (void); // sve vector length in bytes
98+
GGML_BACKEND_API int ggml_cpu_has_sme (void);
9899
// other
99100
GGML_BACKEND_API int ggml_cpu_has_riscv_v (void);
100101
GGML_BACKEND_API int ggml_cpu_has_vsx (void);

‎ggml/src/ggml-cpu/CMakeLists.txt‎

Lines changed: 96 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -111,21 +111,23 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
111111
function(check_arm_feature tag code)
112112
set(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS})
113113
set(CMAKE_REQUIRED_FLAGS "${ARM_MCPU_FLAG}+${tag}")
114-
check_cxx_source_runs(
115-
"${code}"
116-
GGML_MACHINE_SUPPORTS_${tag}
117-
)
114+
check_cxx_source_runs("${code}" GGML_MACHINE_SUPPORTS_${tag})
118115
if (GGML_MACHINE_SUPPORTS_${tag})
119116
set(ARM_MCPU_FLAG_FIX "${ARM_MCPU_FLAG_FIX}+${tag}" PARENT_SCOPE)
120117
else()
121-
set(ARM_MCPU_FLAG_FIX "${ARM_MCPU_FLAG_FIX}+no${tag}" PARENT_SCOPE)
118+
set(CMAKE_REQUIRED_FLAGS "${ARM_MCPU_FLAG}+no${tag}")
119+
check_cxx_source_compiles("int main() { return 0; }" GGML_MACHINE_SUPPORTS_no${tag})
120+
if (GGML_MACHINE_SUPPORTS_no${tag})
121+
set(ARM_MCPU_FLAG_FIX "${ARM_MCPU_FLAG_FIX}+no${tag}" PARENT_SCOPE)
122+
endif()
122123
endif()
123124
set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE})
124125
endfunction()
125126

126127
check_arm_feature(dotprod "#include <arm_neon.h>\nint main() { int8x16_t _a, _b; volatile int32x4_t _s = vdotq_s32(_s, _a, _b); return 0; }")
127128
check_arm_feature(i8mm "#include <arm_neon.h>\nint main() { int8x16_t _a, _b; volatile int32x4_t _s = vmmlaq_s32(_s, _a, _b); return 0; }")
128129
check_arm_feature(sve "#include <arm_sve.h>\nint main() { svfloat32_t _a, _b; volatile svfloat32_t _c = svadd_f32_z(svptrue_b8(), _a, _b); return 0; }")
130+
check_arm_feature(sme "#include <arm_sme.h>\n__arm_locally_streaming int main() { __asm__ volatile(\"smstart; smstop;\"); return 0; }")
129131

130132
list(APPEND ARCH_FLAGS "${ARM_MCPU_FLAG}${ARM_MCPU_FLAG_FIX}")
131133
else()
@@ -150,7 +152,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
150152
if (ARM_FEATURE_RESULT)
151153
message(WARNING "Failed to get ARM features")
152154
else()
153-
foreach(feature DOTPROD SVE MATMUL_INT8 FMA FP16_VECTOR_ARITHMETIC)
155+
foreach(feature DOTPROD SVE MATMUL_INT8 FMA FP16_VECTOR_ARITHMETIC SME)
154156
string(FIND "${ARM_FEATURE}" "__ARM_FEATURE_${feature} 1" feature_pos)
155157
if (NOT ${feature_pos} EQUAL -1)
156158
message(STATUS "ARM feature ${feature} enabled")
@@ -316,6 +318,94 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
316318
target_compile_definitions(${GGML_CPU_NAME} PRIVATE GGML_USE_CPU_AARCH64)
317319
endif()
318320

321+
if (GGML_CPU_KLEIDIAI)
322+
message(STATUS "Using KleidiAI optimized kernels if applicable")
323+
324+
# Disable the KleidiAI tests
325+
set(KLEIDIAI_BUILD_TESTS OFF)
326+
327+
# Fetch KleidiAI sources:
328+
include(FetchContent)
329+
set(KLEIDIAI_COMMIT_TAG "v1.3.0")
330+
set(KLEIDIAI_DOWNLOAD_URL "https://github.com/ARM-software/kleidiai/archive/refs/tags/${KLEIDIAI_COMMIT_TAG}.tar.gz")
331+
set(KLEIDIAI_ARCHIVE_MD5 "060bd2dc64642b091f461cc8dd7426d9")
332+
333+
if (POLICY CMP0135)
334+
cmake_policy(SET CMP0135 NEW)
335+
endif()
336+
337+
FetchContent_Declare(KleidiAI_Download
338+
URL ${KLEIDIAI_DOWNLOAD_URL}
339+
DOWNLOAD_EXTRACT_TIMESTAMP NEW
340+
URL_HASH MD5=${KLEIDIAI_ARCHIVE_MD5})
341+
342+
FetchContent_MakeAvailable(KleidiAI_Download)
343+
FetchContent_GetProperties(KleidiAI_Download
344+
SOURCE_DIR KLEIDIAI_SRC
345+
POPULATED KLEIDIAI_POPULATED)
346+
347+
if (NOT KLEIDIAI_POPULATED)
348+
message(FATAL_ERROR "KleidiAI source downloaded failed.")
349+
endif()
350+
351+
add_compile_definitions(GGML_USE_CPU_KLEIDIAI)
352+
353+
# Remove kleidiai target after fetching it
354+
if (TARGET kleidiai)
355+
set_target_properties(kleidiai PROPERTIES EXCLUDE_FROM_ALL TRUE)
356+
endif()
357+
358+
list(APPEND GGML_CPU_SOURCES
359+
ggml-cpu/kleidiai/kleidiai.cpp
360+
ggml-cpu/kleidiai/kernels.cpp
361+
ggml-cpu/kleidiai/kleidiai.h
362+
ggml-cpu/kleidiai/kernels.h
363+
)
364+
365+
# KleidiAI
366+
include_directories(
367+
${KLEIDIAI_SRC}/
368+
${KLEIDIAI_SRC}/kai/
369+
${KLEIDIAI_SRC}/kai/ukernels/
370+
${KLEIDIAI_SRC}/kai/ukernels/matmul/
371+
${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/
372+
${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/)
373+
374+
set(ARCH_FLAGS_TEMP "${ARCH_FLAGS}")
375+
if (NOT ARCH_FLAGS_TEMP)
376+
string(REGEX MATCH "-march=[^ ]+" ARCH_FLAGS_TEMP "${CMAKE_C_FLAGS}")
377+
endif()
378+
string(FIND "${ARCH_FLAGS_TEMP}" "+dotprod" DOTPROD_ENABLED)
379+
string(FIND "${ARCH_FLAGS_TEMP}" "+i8mm" I8MM_ENABLED)
380+
string(FIND "${ARCH_FLAGS_TEMP}" "+sme" SME_ENABLED)
381+
382+
set(PRIVATE_ARCH_FLAGS ${ARCH_FLAGS})
383+
384+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32.c)
385+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.c)
386+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32p_f32_neon.c)
387+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.c)
388+
389+
if (NOT DOTPROD_ENABLED MATCHES -1)
390+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p4x8_1x4x32_neon_dotprod.c)
391+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4x4_1x4_neon_dotprod.c)
392+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x4_qsi4c32p4x4_16x4_neon_dotprod.c)
393+
endif()
394+
395+
if (NOT I8MM_ENABLED MATCHES -1)
396+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p4x8_16x4_neon_i8mm.c)
397+
endif()
398+
399+
if (NOT SME_ENABLED MATCHES -1)
400+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1vlx4_qsi4c32p4vlx4_1vlx4vl_sme2_mopa.c)
401+
list(APPEND GGML_KLEIDIAI_SOURCES ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qsi4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qsi4c32p4vlx4_1x4vl_sme2_sdot.c)
402+
set(PRIVATE_ARCH_FLAGS "${PRIVATE_ARCH_FLAGS}+sve+sve2")
403+
endif()
404+
405+
set_source_files_properties(${GGML_KLEIDIAI_SOURCES} PROPERTIES COMPILE_OPTIONS "${PRIVATE_ARCH_FLAGS}")
406+
list(APPEND GGML_CPU_SOURCES ${GGML_KLEIDIAI_SOURCES})
407+
endif()
408+
319409
message(STATUS "Adding CPU backend variant ${GGML_CPU_NAME}: ${ARCH_FLAGS} ${ARCH_DEFINITIONS}")
320410
target_sources(${GGML_CPU_NAME} PRIVATE ${GGML_CPU_SOURCES})
321411
target_compile_options(${GGML_CPU_NAME} PRIVATE ${ARCH_FLAGS})

‎ggml/src/ggml-cpu/ggml-cpu.c‎

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,8 @@ struct ggml_arm_arch_features_type {
112112
int has_i8mm;
113113
int has_sve;
114114
int sve_cnt;
115-
} ggml_arm_arch_features = {-1, -1, -1, -1, 0};
115+
int has_sme;
116+
} ggml_arm_arch_features = {-1, -1, -1, -1, 0, -1};
116117
#endif
117118

118119

@@ -2381,15 +2382,20 @@ bool ggml_is_numa(void) {
23812382
#define HWCAP2_I8MM (1 << 13)
23822383
#endif
23832384

2385+
#if !defined(HWCAP2_SME)
2386+
#define HWCAP2_SME (1 << 23)
2387+
#endif
2388+
23842389
static void ggml_init_arm_arch_features(void) {
23852390
#if defined(__linux__) && defined(__aarch64__)
23862391
uint32_t hwcap = getauxval(AT_HWCAP);
23872392
uint32_t hwcap2 = getauxval(AT_HWCAP2);
23882393

2389-
ggml_arm_arch_features.has_neon = !!(hwcap & HWCAP_ASIMD);
2394+
ggml_arm_arch_features.has_neon = !!(hwcap & HWCAP_ASIMD);
23902395
ggml_arm_arch_features.has_dotprod = !!(hwcap & HWCAP_ASIMDDP);
2391-
ggml_arm_arch_features.has_i8mm = !!(hwcap2 & HWCAP2_I8MM);
2392-
ggml_arm_arch_features.has_sve = !!(hwcap & HWCAP_SVE);
2396+
ggml_arm_arch_features.has_i8mm = !!(hwcap2 & HWCAP2_I8MM);
2397+
ggml_arm_arch_features.has_sve = !!(hwcap & HWCAP_SVE);
2398+
ggml_arm_arch_features.has_sme = !!(hwcap2 & HWCAP2_SME);
23932399

23942400
#if defined(__ARM_FEATURE_SVE)
23952401
ggml_arm_arch_features.sve_cnt = PR_SVE_VL_LEN_MASK & prctl(PR_SVE_GET_VL);
@@ -2412,6 +2418,11 @@ static void ggml_init_arm_arch_features(void) {
24122418
}
24132419
ggml_arm_arch_features.has_i8mm = oldp;
24142420

2421+
if (sysctlbyname("hw.optional.arm.FEAT_SME", &oldp, &size, NULL, 0) != 0) {
2422+
oldp = 0;
2423+
}
2424+
ggml_arm_arch_features.has_sme = oldp;
2425+
24152426
ggml_arm_arch_features.has_sve = 0;
24162427
ggml_arm_arch_features.sve_cnt = 0;
24172428
#else
@@ -2435,6 +2446,12 @@ static void ggml_init_arm_arch_features(void) {
24352446
ggml_arm_arch_features.has_sve = 0;
24362447
ggml_arm_arch_features.sve_cnt = 0;
24372448
#endif
2449+
2450+
#if defined(__ARM_FEATURE_SME) || defined(__ARM_FEATURE_SME2)
2451+
ggml_arm_arch_features.has_sme = 1;
2452+
#else
2453+
ggml_arm_arch_features.has_sme = 0;
2454+
#endif
24382455
#endif
24392456
}
24402457
#endif
@@ -14442,6 +14459,14 @@ int ggml_cpu_get_sve_cnt(void) {
1444214459
#endif
1444314460
}
1444414461

14462+
int ggml_cpu_has_sme(void) {
14463+
#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SME)
14464+
return ggml_arm_arch_features.has_sme;
14465+
#else
14466+
return 0;
14467+
#endif
14468+
}
14469+
1444514470
void ggml_cpu_init(void) {
1444614471
// needed to initialize f16 tables
1444714472
{

‎ggml/src/ggml-cpu/ggml-cpu.cpp‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@
1414
#include "ggml-cpu-hbm.h"
1515
#endif
1616

17+
#ifdef GGML_USE_CPU_KLEIDIAI
18+
#include "kleidiai/kleidiai.h"
19+
#endif
20+
1721
#if defined(__APPLE__)
1822
#include <sys/types.h>
1923
#include <sys/sysctl.h>
@@ -39,6 +43,12 @@ std::vector<ggml_backend_buffer_type_t>& ggml_backend_cpu_get_extra_buffers_type
3943
}
4044
#endif
4145

46+
#ifdef GGML_USE_CPU_KLEIDIAI
47+
if (ggml_backend_cpu_kleidiai_buffer_type()) {
48+
bufts.push_back(ggml_backend_cpu_kleidiai_buffer_type());
49+
}
50+
#endif
51+
4252
#ifdef GGML_USE_CPU_AARCH64
4353
if (ggml_backend_cpu_aarch64_buffer_type()) {
4454
bufts.push_back(ggml_backend_cpu_aarch64_buffer_type());
@@ -538,6 +548,9 @@ static ggml_backend_feature * ggml_backend_cpu_get_features(ggml_backend_reg_t r
538548
static std::string sve_cnt = std::to_string(ggml_cpu_get_sve_cnt());
539549
features.push_back({ "SVE_CNT", sve_cnt.c_str() });
540550
}
551+
if (ggml_cpu_has_sme()) {
552+
features.push_back({ "SME", "1" });
553+
}
541554
if (ggml_cpu_has_riscv_v()) {
542555
features.push_back({ "RISCV_V", "1" });
543556
}
@@ -559,6 +572,9 @@ static ggml_backend_feature * ggml_backend_cpu_get_features(ggml_backend_reg_t r
559572
#ifdef GGML_USE_OPENMP
560573
features.push_back({ "OPENMP", "1" });
561574
#endif
575+
#ifdef GGML_USE_CPU_KLEIDIAI
576+
features.push_back({ "KLEIDIAI", "1" });
577+
#endif
562578
#ifdef GGML_USE_CPU_AARCH64
563579
features.push_back({ "AARCH64_REPACK", "1" });
564580
#endif

0 commit comments

Comments
 (0)