Skip to content

Commit eb29b03

Browse files
johnnynunezjohnnynunez
authored
Update sgemm.cpp
1 parent e63e3dc commit eb29b03

1 file changed

Lines changed: 144 additions & 71 deletions

File tree

‎llamafile/sgemm.cpp‎

Lines changed: 144 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,8 @@
1+
// Adapted from
2+
// https://github.com/Mozilla-Ocho/llamafile/blob/0.8.8/llamafile/sgemm.cpp
3+
// Copyrigth 2024 Mozilla Foundation.
4+
// Copyright(c) 2024 by KVCache.AI, All Rights Reserved.
5+
16
// -*- mode:c++;indent-tabs-mode:nil;c-basic-offset:4;coding:utf-8 -*-
27
// vi: set et ft=cpp ts=4 sts=4 sw=4 fenc=utf-8 :vi
38
//
@@ -16,75 +21,138 @@
1621
// limitations under the License.
1722

1823
#include "sgemm.h"
19-
#include "llamafile.h"
24+
// #include <cosmo.h>
25+
// #include <cpuid.h>
26+
// #include <libc/sysv/consts/hwcap.h>
27+
#include <stdio.h>
2028
#include <cassert>
21-
#include <cosmo.h>
22-
#include <cpuid.h>
23-
#include <libc/sysv/consts/hwcap.h>
29+
// #include "llamafile.h"
30+
31+
// ARM64-specific headers and constants
32+
#ifdef __aarch64__
2433
#include <sys/auxv.h>
34+
#include <asm/hwcap.h>
35+
#endif
2536

2637
static const struct GemmFuncs {
27-
typeof(llamafile_sgemm) *sgemm;
28-
typeof(llamafile_mixmul) *mixmul;
29-
typeof(llamafile_mixmul_iqk) *iqk_mixmul = iqk_mul_mat_moe_unsupported;
38+
bool (*sgemm)(long, long, long, const void*, long, const void*, long, void*, long, int, int, int, int, int, int, int);
39+
bool (*mixmul)(const struct ggml_compute_params*, const struct ggml_tensor*, const struct ggml_tensor*, const struct ggml_tensor*, struct ggml_tensor*);
40+
bool (*iqk_mixmul)(long, long, long, int, int, const void*, const void*, float*, long, long, const void*, int, int);
41+
// typeof(llamafile_sgemm)* sgemm;
42+
// typeof(llamafile_mixmul)* mixmul;
43+
// typeof(llamafile_mixmul_iqk)* iqk_mixmul = iqk_mul_mat_moe_unsupported;
3044
GemmFuncs() {
31-
#ifdef __x86_64__
32-
if (X86_HAVE(AVX)) {
33-
if (X86_HAVE(FMA)) {
34-
if (X86_HAVE(AVX2)) {
35-
if (X86_HAVE(AVX512F)) {
36-
if (X86_HAVE(AVX512VL) && //
37-
X86_HAVE(AVX512BW) && //
38-
X86_HAVE(AVX512DQ) && //
39-
X86_HAVE(AVX512_VNNI) && //
40-
X86_HAVE(AVX512_BF16)) {
41-
// AMD Zen4+ (2023-)
42-
sgemm = llamafile_sgemm_amd_zen4;
43-
mixmul = llamafile_mixmul_amd_zen4;
44-
iqk_mixmul = iqk_mul_mat_moe_zen4;
45-
} else {
46-
// Intel Xeon Skylake+ (2015-)
47-
sgemm = llamafile_sgemm_amd_avx512f;
48-
mixmul = llamafile_mixmul_amd_avx512f;
49-
iqk_mixmul = iqk_mul_mat_moe;
50-
}
51-
} else if (X86_HAVE(AVXVNNI)) {
52-
// Intel Alderlake (2021-)
53-
sgemm = llamafile_sgemm_amd_avxvnni;
54-
mixmul = llamafile_mixmul_amd_avxvnni;
55-
iqk_mixmul = iqk_mul_mat_moe;
56-
} else {
57-
// Intel Haswell/Broadwell/Skylake (2013-2020)
58-
// AMD Excavator (2015-2022)
59-
sgemm = llamafile_sgemm_amd_avx2;
60-
mixmul = llamafile_mixmul_amd_avx2;
61-
if (X86_HAVE(F16C))
62-
iqk_mixmul = iqk_mul_mat_moe;
63-
}
64-
} else {
65-
// AMD Piledriver (2011-2014)
66-
sgemm = llamafile_sgemm_amd_fma;
67-
mixmul = llamafile_mixmul_amd_fma;
68-
if (X86_HAVE(F16C))
69-
iqk_mixmul = iqk_mul_mat_moe;
70-
}
71-
} else {
72-
// Intel Sandybridge/Ivybridge (2010-2012)
73-
// AMD Bulldozer (2011)
74-
sgemm = llamafile_sgemm_amd_avx;
75-
mixmul = llamafile_mixmul_amd_avx;
76-
}
77-
} else {
78-
// AMD K8/Barcelona (2003-2010)
79-
// Intel Core/Nehalem (2006-2009)
80-
sgemm = llamafile_sgemm_unsupported;
81-
mixmul = llamafile_mixmul_unsupported;
82-
}
45+
#if defined(__x86_64__) || defined(_M_X64)
46+
// if (X86_HAVE(AVX)) {
47+
// if (X86_HAVE(FMA)) {
48+
// if (X86_HAVE(AVX2)) {
49+
// if (X86_HAVE(AVX512F)) {
50+
// if (X86_HAVE(AVX512VL) && //
51+
// X86_HAVE(AVX512BW) && //
52+
// X86_HAVE(AVX512DQ) && //
53+
// X86_HAVE(AVX512_VNNI) && //
54+
// X86_HAVE(AVX512_BF16)) {
55+
// // AMD Zen4+ (2023-)
56+
// sgemm = llamafile_sgemm_amd_zen4;
57+
// mixmul = llamafile_mixmul_amd_zen4;
58+
// iqk_mixmul = iqk_mul_mat_moe_zen4;
59+
// } else {
60+
// // Intel Xeon Skylake+ (2015-)
61+
// sgemm = llamafile_sgemm_amd_avx512f;
62+
// mixmul = llamafile_mixmul_amd_avx512f;
63+
// iqk_mixmul = iqk_mul_mat_moe;
64+
// }
65+
// } else if (X86_HAVE(AVXVNNI)) {
66+
// // Intel Alderlake (2021-)
67+
// sgemm = llamafile_sgemm_amd_avxvnni;
68+
// mixmul = llamafile_mixmul_amd_avxvnni;
69+
// iqk_mixmul = iqk_mul_mat_moe;
70+
// } else {
71+
// // Intel Haswell/Broadwell/Skylake (2013-2020)
72+
// // AMD Excavator (2015-2022)
73+
// sgemm = llamafile_sgemm_amd_avx2;
74+
// mixmul = llamafile_mixmul_amd_avx2;
75+
// if (X86_HAVE(F16C))
76+
// iqk_mixmul = iqk_mul_mat_moe;
77+
// }
78+
// } else {
79+
// // AMD Piledriver (2011-2014)
80+
// sgemm = llamafile_sgemm_amd_fma;
81+
// mixmul = llamafile_mixmul_amd_fma;
82+
// if (X86_HAVE(F16C))
83+
// iqk_mixmul = iqk_mul_mat_moe;
84+
// }
85+
// } else {
86+
// // Intel Sandybridge/Ivybridge (2010-2012)
87+
// // AMD Bulldozer (2011)
88+
// sgemm = llamafile_sgemm_amd_avx;
89+
// mixmul = llamafile_mixmul_amd_avx;
90+
// }
91+
// } else {
92+
// // AMD K8/Barcelona (2003-2010)
93+
// // Intel Core/Nehalem (2006-2009)
94+
// sgemm = llamafile_sgemm_unsupported;
95+
// mixmul = llamafile_mixmul_unsupported;
96+
// }
97+
98+
#if defined(__AVX__)
99+
#if defined(__FMA__) || (defined(_MSC_VER) && (defined(__AVX2__) || defined(__AVX512F__)))
100+
#if defined(__AVX2__)
101+
#if defined(__AVX512F__)
102+
#if defined(__AVX512VL__) && defined(__AVX512BW__) && defined(__AVX512DQ__) && defined(__AVX512VNNI__) && defined(__AVX512BF16__)
103+
// AMD Zen4+ (2023-)
104+
sgemm = llamafile_sgemm_amd_zen4;
105+
mixmul = llamafile_mixmul_amd_zen4;
106+
iqk_mixmul = iqk_mul_mat_moe_zen4;
107+
#else
108+
// Intel Xeon Skylake+ (2015-)
109+
sgemm = llamafile_sgemm_amd_avx512f;
110+
mixmul = llamafile_mixmul_amd_avx512f;
111+
iqk_mixmul = iqk_mul_mat_moe;
112+
#endif
113+
#elif defined(__AVXVNNI__)
114+
// Intel Alderlake (2021-)
115+
sgemm = llamafile_sgemm_amd_avxvnni;
116+
mixmul = llamafile_mixmul_amd_avxvnni;
117+
iqk_mixmul = iqk_mul_mat_moe;
118+
#else
119+
// Intel Haswell/Broadwell/Skylake (2013-2020)
120+
// AMD Excavator (2015-2022)
121+
sgemm = llamafile_sgemm_amd_avx2;
122+
mixmul = llamafile_mixmul_amd_avx2;
123+
#if defined(__F16C__)
124+
iqk_mixmul = iqk_mul_mat_moe;
125+
#endif
126+
#endif
127+
#else
128+
// AMD Piledriver (2011-2014)
129+
sgemm = llamafile_sgemm_amd_fma;
130+
mixmul = llamafile_mixmul_amd_fma;
131+
#if defined(__F16C__)
132+
iqk_mixmul = iqk_mul_mat_moe;
133+
#endif
134+
#endif
135+
#else
136+
// Intel Sandybridge/Ivybridge (2010-2012)
137+
// AMD Bulldozer (2011)
138+
sgemm = llamafile_sgemm_amd_avx;
139+
mixmul = llamafile_mixmul_amd_avx;
140+
#endif
141+
#else
142+
// AMD K8/Barcelona (2003-2010)
143+
// Intel Core/Nehalem (2006-2009)
144+
sgemm = llamafile_sgemm_unsupported;
145+
mixmul = llamafile_mixmul_unsupported;
146+
#endif
147+
83148
#elif defined(__aarch64__)
84-
long hwcap = getauxval(AT_HWCAP);
85-
if ((hwcap & HWCAP_FPHP) && // fp16 scalar isa (ID_AA64PFR0_EL1.FP == 1)
86-
(hwcap & HWCAP_ASIMDHP) && // fp16 vector isa (ID_AA64PFR0_EL1.AdvSIMD == 1)
87-
(hwcap & HWCAP_ASIMDDP)) { // dotprod isa (ID_AA64ISAR0_EL1.DP == 1)
149+
// ARM64 hardware capability detection
150+
long hwcap = 0;
151+
#ifdef __linux__
152+
hwcap = getauxval(AT_HWCAP);
153+
if ((hwcap & HWCAP_FPHP) && // fp16 scalar isa (ID_AA64PFR0_EL1.FP == 1)
154+
(hwcap & HWCAP_ASIMDHP) && // fp16 vector isa (ID_AA64PFR0_EL1.AdvSIMD == 1)
155+
(hwcap & HWCAP_ASIMDDP)) { // dotprod isa (ID_AA64ISAR0_EL1.DP == 1)
88156
// e.g. Apple M1, Raspberry Pi 5
89157
sgemm = llamafile_sgemm_arm82;
90158
mixmul = llamafile_mixmul_arm82;
@@ -94,6 +162,12 @@ static const struct GemmFuncs {
94162
sgemm = llamafile_sgemm_arm80;
95163
mixmul = llamafile_mixmul_arm80;
96164
}
165+
#else
166+
// Non-Linux ARM64 systems (e.g., macOS)
167+
// Use baseline ARM64 implementation
168+
sgemm = llamafile_sgemm_arm80;
169+
mixmul = llamafile_mixmul_arm80;
170+
#endif
97171
#else
98172
sgemm = llamafile_sgemm_unsupported;
99173
mixmul = llamafile_mixmul_unsupported;
@@ -120,26 +194,25 @@ static const struct GemmFuncs {
120194
* @param ldc is row stride of `C`
121195
* @param ith is thread id (must be less than `nth`)
122196
* @param nth is number of threads (must be greater than zero)
197+
* @param task is GGML task type
123198
* @param Atype is GGML data type of `A`
124199
* @param Btype is GGML data type of `B`
125200
* @param Ctype is GGML data type of `C`
201+
* @param precision may be used to control the internal compute type
126202
* @return true if this function was able to service the matmul request
127203
*/
128-
bool llamafile_sgemm(long m, long n, long k, const void *A, long lda, const void *B, long ldb,
129-
void *C, long ldc, int ith, int nth, int Atype, int Btype, int Ctype) {
130-
return funcs.sgemm(m, n, k, A, lda, B, ldb, C, ldc, ith, nth, Atype, Btype, Ctype);
204+
bool llamafile_sgemm(long m, long n, long k, const void* A, long lda, const void* B, long ldb, void* C, long ldc, int ith, int nth, int task, int Atype, int Btype, int Ctype, int precision) {
205+
return funcs.sgemm(m, n, k, A, lda, B, ldb, C, ldc, ith, nth, task, Atype, Btype, Ctype,
206+
precision);
131207
}
132208

133209
/**
134210
* Performs "mixture of experts" tensor multiplication on CPU.
135211
*/
136-
bool llamafile_mixmul(const ggml_compute_params *params, const ggml_tensor *weights,
137-
const ggml_tensor *thought, const ggml_tensor *plan, ggml_tensor *result) {
212+
bool llamafile_mixmul(const ggml_compute_params* params, const ggml_tensor* weights, const ggml_tensor* thought, const ggml_tensor* plan, ggml_tensor* result) {
138213
return funcs.mixmul(params, weights, thought, plan, result);
139214
}
140215

141-
bool llamafile_mixmul_iqk(long Nx, long Ny, long ne00, int ne11, int typeA, const void *A,
142-
const void *B, float *C, long nb1, long nb2, const void *vrow_mapping,
143-
int ith, int nth) {
216+
bool llamafile_mixmul_iqk(long Nx, long Ny, long ne00, int ne11, int typeA, const void* A, const void* B, float* C, long nb1, long nb2, const void* vrow_mapping, int ith, int nth) {
144217
return funcs.iqk_mixmul(Nx, Ny, ne00, ne11, typeA, A, B, C, nb1, nb2, vrow_mapping, ith, nth);
145218
}

0 commit comments

Comments
 (0)