char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Jonathan Clohessykleidiai: Rework KleidiAI Build System/Integration (#26077)d222767c7

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
3.3 KiB101 linesraw
1// SPDX-FileCopyrightText: Copyright 2025-2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
2// SPDX-License-Identifier: MIT
3//
4
5#pragma once
6
7#include "ggml.h"
8
9enum cpu_feature {
10    CPU_FEATURE_NONE    = 0,
11    CPU_FEATURE_DOTPROD = 1,
12    CPU_FEATURE_I8MM    = 2,
13    CPU_FEATURE_SVE     = 4,
14    CPU_FEATURE_SME     = 8,
15    CPU_FEATURE_SME2    = 16,
16    CPU_FEATURE_FP16    = 32
17};
18
19inline cpu_feature& operator|=(cpu_feature& lhs, cpu_feature rhs) {
20    lhs = static_cast<cpu_feature>(lhs | rhs);
21    return lhs;
22}
23inline cpu_feature operator|(cpu_feature lhs, cpu_feature rhs) {
24    return static_cast<cpu_feature>(static_cast<int>(lhs) | static_cast<int>(rhs));
25}
26
27struct kernel_info {
28    size_t (*get_m_step)(void);
29    size_t (*get_n_step)(void);
30    size_t (*get_mr)(void);
31    size_t (*get_nr)(void);
32    size_t (*get_kr)(void);
33    size_t (*get_sr)(void);
34
35    size_t (*get_dst_offset)(size_t m_idx, size_t n_idx, size_t stride);
36    size_t (*get_dst_size)(size_t m, size_t n);
37
38    size_t (*get_lhs_offset_ex)(size_t m_idx, size_t k, size_t bl);
39
40    size_t (*get_rhs_packed_offset_ex)(size_t n_idx, size_t k, size_t bl);
41
42    void (*run_kernel_ex)(
43        size_t m, size_t n, size_t k, size_t bl,
44        const void* lhs_packed, const void* rhs_packed,
45        void* dst, size_t dst_stride_row, size_t dst_stride_col,
46        float clamp_min, float clamp_max);
47};
48
49struct lhs_packing_info {
50    size_t (*get_offset)(size_t m_idx, size_t lhs_stride);
51
52    size_t (*get_packed_offset_ex)(size_t m_idx, size_t k, size_t bl, size_t mr, size_t kr, size_t sr);
53
54    size_t (*packed_size_ex)(size_t m, size_t k, size_t bl, size_t mr, size_t kr, size_t sr);
55
56    void (*pack_func_ex)(size_t m, size_t k, size_t bl, size_t mr, size_t kr, size_t sr,
57        size_t m_idx_start, const void * lhs, size_t lhs_stride, void * lhs_packed);
58};
59
60enum rhs_repack_mode {
61    RHS_REPACK_PER_KERNEL,
62    RHS_REPACK_SHARED,
63    RHS_REPACK_SINGLE_ONLY,
64};
65
66struct rhs_packing_info {
67    size_t (*packed_stride)(size_t k, size_t nr, size_t kr, size_t bl);
68
69    void (*to_float)(const void *packed_data, int32_t row_idx, int64_t nc, float *out,
70                     size_t nr_pack, size_t packed_row_stride, size_t kr, size_t bl,
71                     size_t num_bytes_multiplier);
72
73    size_t (*packed_size_ex)(size_t n, size_t k, size_t nr, size_t kr, size_t bl);
74
75    size_t (*packed_stride_ex)(size_t k, size_t nr, size_t kr, size_t bl);
76
77    void (*pack_func_ex)(size_t num_groups, size_t n, size_t k, size_t nr, size_t kr, size_t sr, size_t bl,
78        size_t rhs_stride, const void * rhs, const void * bias, const void * scale, void * rhs_packed, size_t extra_bytes, const void * params);
79
80    rhs_repack_mode repack_mode = RHS_REPACK_PER_KERNEL;
81};
82
83struct ggml_kleidiai_kernels {
84    kernel_info      gemm;
85    lhs_packing_info gemm_lhs_info;
86
87    kernel_info      gemv;
88    lhs_packing_info gemv_lhs_info;
89
90    rhs_packing_info rhs_info;
91
92    cpu_feature required_cpu;
93    ggml_type lhs_type;
94    ggml_type rhs_type;
95    ggml_type op_type;
96};
97
98ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, const ggml_tensor * tensor);
99ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features);
100ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features);
101ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_f32(cpu_feature features);