sgemm.cpp 40 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148
  1. // -*- mode:c++;indent-tabs-mode:nil;c-basic-offset:4;coding:utf-8 -*-
  2. // vi: set et ft=c++ ts=4 sts=4 sw=4 fenc=utf-8 :vi
  3. //
  4. // Copyright 2024 Mozilla Foundation
  5. //
  6. // Permission is hereby granted, free of charge, to any person obtaining
  7. // a copy of this software and associated documentation files (the
  8. // "Software"), to deal in the Software without restriction, including
  9. // without limitation the rights to use, copy, modify, merge, publish,
  10. // distribute, sublicense, and/or sell copies of the Software, and to
  11. // permit persons to whom the Software is furnished to do so, subject to
  12. // the following conditions:
  13. //
  14. // The above copyright notice and this permission notice shall be
  15. // included in all copies or substantial portions of the Software.
  16. //
  17. // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  18. // EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  19. // MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  20. // NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
  21. // BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
  22. // ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
  23. // CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  24. // SOFTWARE.
  25. //
  26. // _ _ ___ _ _ ___
  27. // | |_(_)_ _ _ _| _ ) | /_\ / __|
  28. // | _| | ' \ || | _ \ |__ / _ \\__ \.
  29. // \__|_|_||_\_, |___/____/_/ \_\___/
  30. // |__/
  31. //
  32. // BASIC LINEAR ALGEBRA SUBPROGRAMS
  33. //
  34. //
  35. // This file implements multithreaded CPU matrix multiplication for the
  36. // common contiguous use case C = Aᵀ * B. These kernels are designed to
  37. // have excellent performance[1] for matrices that fit in the CPU cache
  38. // without imposing any overhead such as cache filling or malloc calls.
  39. //
  40. // This implementation does not guarantee any upper bound with rounding
  41. // errors, which grow along with k. Our goal's to maximally exploit the
  42. // hardware for performance, and then use whatever resources remain for
  43. // improving numerical accuracy.
  44. //
  45. // [1] J. Tunney, ‘LLaMA Now Goes Faster on CPUs’, Mar. 2024. [Online].
  46. // Available: https://justine.lol/matmul/. [Accessed: 29-Mar-2024].
  47. #pragma GCC diagnostic ignored "-Wpedantic"
  48. #pragma GCC diagnostic ignored "-Wignored-attributes"
  49. #include "sgemm.h"
  50. #include "ggml-impl.h"
  51. #include "ggml-quants.h"
  52. #ifdef _MSC_VER
  53. #define NOINLINE __declspec(noinline)
  54. #else
  55. #define NOINLINE __attribute__((__noinline__))
  56. #endif
  57. #if defined(__ARM_NEON) || defined(__AVX512F__)
  58. #define VECTOR_REGISTERS 32
  59. #else
  60. #define VECTOR_REGISTERS 16
  61. #endif
  62. // there will be blocks
  63. #define BEGIN_KERNEL(RM, RN) \
  64. int ytiles = (m - m0) / RM; \
  65. int xtiles = (n - n0) / RN; \
  66. int tiles = ytiles * xtiles; \
  67. int duty = (tiles + nth - 1) / nth; \
  68. int start = duty * ith; \
  69. int end = start + duty; \
  70. if (end > tiles) \
  71. end = tiles; \
  72. for (int job = start; job < end; ++job) { \
  73. int i = m0 + job / xtiles * RM; \
  74. int j = n0 + job % xtiles * RN;
  75. #define END_KERNEL() }
  76. #define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1)
  77. namespace {
  78. inline float unhalf(ggml_fp16_t d) {
  79. return GGML_FP16_TO_FP32(d);
  80. }
  81. ////////////////////////////////////////////////////////////////////////////////////////////////////
  82. // VECTORIZED ARITHMETIC OPERATIONS
  83. #if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
  84. inline __m128 add(__m128 x, __m128 y) { return _mm_add_ps(x, y); }
  85. inline __m128 sub(__m128 x, __m128 y) { return _mm_sub_ps(x, y); }
  86. inline __m128 mul(__m128 x, __m128 y) { return _mm_mul_ps(x, y); }
  87. #endif // __SSE__
  88. #if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
  89. inline __m256 add(__m256 x, __m256 y) { return _mm256_add_ps(x, y); }
  90. inline __m256 sub(__m256 x, __m256 y) { return _mm256_sub_ps(x, y); }
  91. inline __m256 mul(__m256 x, __m256 y) { return _mm256_mul_ps(x, y); }
  92. #endif // __AVX__
  93. #if defined(__AVX512F__)
  94. inline __m512 add(__m512 x, __m512 y) { return _mm512_add_ps(x, y); }
  95. inline __m512 sub(__m512 x, __m512 y) { return _mm512_sub_ps(x, y); }
  96. inline __m512 mul(__m512 x, __m512 y) { return _mm512_mul_ps(x, y); }
  97. #endif // __AVX512F__
  98. #if defined(__ARM_NEON)
  99. inline float32x4_t add(float32x4_t x, float32x4_t y) { return vaddq_f32(x, y); }
  100. inline float32x4_t sub(float32x4_t x, float32x4_t y) { return vsubq_f32(x, y); }
  101. inline float32x4_t mul(float32x4_t x, float32x4_t y) { return vmulq_f32(x, y); }
  102. #endif // __ARM_NEON
  103. #if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC)
  104. inline float16x8_t add(float16x8_t x, float16x8_t y) { return vaddq_f16(x, y); }
  105. inline float16x8_t sub(float16x8_t x, float16x8_t y) { return vsubq_f16(x, y); }
  106. inline float16x8_t mul(float16x8_t x, float16x8_t y) { return vmulq_f16(x, y); }
  107. #endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
  108. ////////////////////////////////////////////////////////////////////////////////////////////////////
  109. // VECTORIZED HORIZONTAL SUM
  110. #if defined(__ARM_NEON)
  111. inline float hsum(float32x4_t x) {
  112. return vaddvq_f32(x);
  113. }
  114. #endif // __ARM_NEON
  115. #if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER)
  116. inline float hsum(float16x8_t x) {
  117. return vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(x)),
  118. vcvt_f32_f16(vget_high_f16(x))));
  119. }
  120. #endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
  121. #if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
  122. inline float hsum(__m128 x) {
  123. #if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
  124. x = _mm_add_ps(x, _mm_movehl_ps(x, x));
  125. x = _mm_add_ss(x, _mm_movehdup_ps(x));
  126. #else
  127. __m128 t;
  128. t = _mm_shuffle_ps(x, x, _MM_SHUFFLE(2, 3, 0, 1));
  129. x = _mm_add_ps(x, t);
  130. t = _mm_movehl_ps(t, x);
  131. x = _mm_add_ss(x, t);
  132. #endif
  133. return _mm_cvtss_f32(x);
  134. }
  135. #endif
  136. #if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
  137. inline float hsum(__m256 x) {
  138. return hsum(_mm_add_ps(_mm256_extractf128_ps(x, 1),
  139. _mm256_castps256_ps128(x)));
  140. }
  141. #endif // __AVX__
  142. #if defined(__AVX512F__)
  143. inline float hsum(__m512 x) {
  144. return _mm512_reduce_add_ps(x);
  145. }
  146. #endif // __AVX512F__
  147. ////////////////////////////////////////////////////////////////////////////////////////////////////
  148. // VECTORIZED MEMORY LOADING
  149. template <typename T, typename U> T load(const U *);
  150. #if defined(__ARM_NEON)
  151. template <> inline float32x4_t load(const float *p) {
  152. return vld1q_f32(p);
  153. }
  154. #if !defined(_MSC_VER)
  155. template <> inline float16x8_t load(const ggml_fp16_t *p) {
  156. return vld1q_f16((const float16_t *)p);
  157. }
  158. template <> inline float32x4_t load(const ggml_fp16_t *p) {
  159. return vcvt_f32_f16(vld1_f16((const float16_t *)p));
  160. }
  161. #endif // _MSC_VER
  162. #endif // __ARM_NEON
  163. #if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
  164. template <> inline __m128 load(const float *p) {
  165. return _mm_loadu_ps(p);
  166. }
  167. #endif // __SSE__
  168. #if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
  169. template <> inline __m256 load(const float *p) {
  170. return _mm256_loadu_ps(p);
  171. }
  172. #endif // __AVX__
  173. #if defined(__F16C__)
  174. template <> inline __m256 load(const ggml_fp16_t *p) {
  175. return _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)p));
  176. }
  177. #endif // __F16C__
  178. #if defined(__AVX512F__)
  179. template <> inline __m512 load(const float *p) {
  180. return _mm512_loadu_ps(p);
  181. }
  182. template <> inline __m512 load(const ggml_fp16_t *p) {
  183. return _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)p));
  184. }
  185. #endif // __AVX512F__
  186. ////////////////////////////////////////////////////////////////////////////////////////////////////
  187. // ABSTRACTIONS
  188. /**
  189. * Computes a * b + c.
  190. *
  191. * This operation will become fused into a single arithmetic instruction
  192. * if the hardware has support for this feature, e.g. Intel Haswell+ (c.
  193. * 2013), AMD Bulldozer+ (c. 2011), etc.
  194. */
  195. template <typename T, typename U>
  196. inline U madd(T a, T b, U c) {
  197. return add(mul(a, b), c);
  198. }
  199. /**
  200. * Computes a * b + c with error correction.
  201. *
  202. * @see W. Kahan, "Further remarks on reducing truncation errors,"
  203. * Communications of the ACM, vol. 8, no. 1, p. 40, Jan. 1965,
  204. * doi: 10.1145/363707.363723.
  205. */
  206. template <typename T, typename U>
  207. inline U madder(T a, T b, U c, U *e) {
  208. U y = sub(mul(a, b), *e);
  209. U t = add(c, y);
  210. *e = sub(sub(t, c), y);
  211. return t;
  212. }
  213. ////////////////////////////////////////////////////////////////////////////////////////////////////
  214. // FLOATING POINT MATRIX MULTIPLICATION
  215. template <int KN, typename D, typename V, typename TA, typename TB, typename TC>
  216. class tinyBLAS {
  217. public:
  218. tinyBLAS(int k,
  219. const TA *A, int lda,
  220. const TB *B, int ldb,
  221. TC *C, int ldc,
  222. int ith, int nth)
  223. : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) {
  224. }
  225. void matmul(int m, int n, int task) {
  226. if (task == GGML_TASK_TYPE_COMPUTE)
  227. mnpack(0, m, 0, n);
  228. }
  229. private:
  230. NOINLINE void mnpack(int m0, int m, int n0, int n) {
  231. int mc, nc, mp, np;
  232. if (m - m0 <= 0 || n - n0 <= 0)
  233. return;
  234. if (VECTOR_REGISTERS >= 32 && n - n0 >= 5 && m - m0 >= 5) {
  235. mc = 5;
  236. nc = 5;
  237. gemm5x5(m0, m, n0, n);
  238. } else if (n - n0 >= 4 && m - m0 >= 3) {
  239. mc = 3;
  240. nc = 4;
  241. gemm3x4(m0, m, n0, n);
  242. } else if (n - n0 >= 4) {
  243. mc = 1;
  244. nc = 4;
  245. gemm1x4(m0, m, n0, n);
  246. } else if (m - m0 >= 4) {
  247. mc = 4;
  248. nc = 1;
  249. gemm4x1(m0, m, n0, n);
  250. } else {
  251. mc = 1;
  252. nc = 1;
  253. gemm1x1(m0, m, n0, n);
  254. }
  255. mp = m0 + (m - m0) / mc * mc;
  256. np = n0 + (n - n0) / nc * nc;
  257. mnpack(mp, m, n0, np);
  258. mnpack(m0, mp, np, n);
  259. mnpack(mp, m, np, n);
  260. }
  261. NOINLINE void gemm5x5(int m0, int m, int n0, int n) {
  262. BEGIN_KERNEL(5, 5)
  263. D c00 = {0};
  264. D c01 = {0};
  265. D c02 = {0};
  266. D c03 = {0};
  267. D c04 = {0};
  268. D c10 = {0};
  269. D c11 = {0};
  270. D c12 = {0};
  271. D c13 = {0};
  272. D c14 = {0};
  273. D c20 = {0};
  274. D c21 = {0};
  275. D c22 = {0};
  276. D c23 = {0};
  277. D c24 = {0};
  278. D c30 = {0};
  279. D c31 = {0};
  280. D c32 = {0};
  281. D c33 = {0};
  282. D c34 = {0};
  283. D c40 = {0};
  284. D c41 = {0};
  285. D c42 = {0};
  286. D c43 = {0};
  287. D c44 = {0};
  288. for (int l = 0; l < k; l += KN) {
  289. V k0 = load<V>(B + ldb * (j + 0) + l);
  290. V k1 = load<V>(B + ldb * (j + 1) + l);
  291. V k2 = load<V>(B + ldb * (j + 2) + l);
  292. V k3 = load<V>(B + ldb * (j + 3) + l);
  293. V k4 = load<V>(B + ldb * (j + 4) + l);
  294. V a0 = load<V>(A + lda * (i + 0) + l);
  295. c00 = madd(a0, k0, c00);
  296. c01 = madd(a0, k1, c01);
  297. c02 = madd(a0, k2, c02);
  298. c03 = madd(a0, k3, c03);
  299. c04 = madd(a0, k4, c04);
  300. V a1 = load<V>(A + lda * (i + 1) + l);
  301. c10 = madd(a1, k0, c10);
  302. c11 = madd(a1, k1, c11);
  303. c12 = madd(a1, k2, c12);
  304. c13 = madd(a1, k3, c13);
  305. c14 = madd(a1, k4, c14);
  306. V a2 = load<V>(A + lda * (i + 2) + l);
  307. c20 = madd(a2, k0, c20);
  308. c21 = madd(a2, k1, c21);
  309. c22 = madd(a2, k2, c22);
  310. c23 = madd(a2, k3, c23);
  311. c24 = madd(a2, k4, c24);
  312. V a3 = load<V>(A + lda * (i + 3) + l);
  313. c30 = madd(a3, k0, c30);
  314. c31 = madd(a3, k1, c31);
  315. c32 = madd(a3, k2, c32);
  316. c33 = madd(a3, k3, c33);
  317. c34 = madd(a3, k4, c34);
  318. V a4 = load<V>(A + lda * (i + 4) + l);
  319. c40 = madd(a4, k0, c40);
  320. c41 = madd(a4, k1, c41);
  321. c42 = madd(a4, k2, c42);
  322. c43 = madd(a4, k3, c43);
  323. c44 = madd(a4, k4, c44);
  324. }
  325. C[ldc * (j + 0) + (i + 0)] = hsum(c00);
  326. C[ldc * (j + 0) + (i + 1)] = hsum(c10);
  327. C[ldc * (j + 0) + (i + 2)] = hsum(c20);
  328. C[ldc * (j + 0) + (i + 3)] = hsum(c30);
  329. C[ldc * (j + 0) + (i + 4)] = hsum(c40);
  330. C[ldc * (j + 1) + (i + 0)] = hsum(c01);
  331. C[ldc * (j + 1) + (i + 1)] = hsum(c11);
  332. C[ldc * (j + 1) + (i + 2)] = hsum(c21);
  333. C[ldc * (j + 1) + (i + 3)] = hsum(c31);
  334. C[ldc * (j + 1) + (i + 4)] = hsum(c41);
  335. C[ldc * (j + 2) + (i + 0)] = hsum(c02);
  336. C[ldc * (j + 2) + (i + 1)] = hsum(c12);
  337. C[ldc * (j + 2) + (i + 2)] = hsum(c22);
  338. C[ldc * (j + 2) + (i + 3)] = hsum(c32);
  339. C[ldc * (j + 2) + (i + 4)] = hsum(c42);
  340. C[ldc * (j + 3) + (i + 0)] = hsum(c03);
  341. C[ldc * (j + 3) + (i + 1)] = hsum(c13);
  342. C[ldc * (j + 3) + (i + 2)] = hsum(c23);
  343. C[ldc * (j + 3) + (i + 3)] = hsum(c33);
  344. C[ldc * (j + 3) + (i + 4)] = hsum(c43);
  345. C[ldc * (j + 4) + (i + 0)] = hsum(c04);
  346. C[ldc * (j + 4) + (i + 1)] = hsum(c14);
  347. C[ldc * (j + 4) + (i + 2)] = hsum(c24);
  348. C[ldc * (j + 4) + (i + 3)] = hsum(c34);
  349. C[ldc * (j + 4) + (i + 4)] = hsum(c44);
  350. END_KERNEL()
  351. }
  352. NOINLINE void gemm3x4(int m0, int m, int n0, int n) {
  353. BEGIN_KERNEL(3, 4)
  354. D c00 = {0};
  355. D c01 = {0};
  356. D c02 = {0};
  357. D c03 = {0};
  358. D c10 = {0};
  359. D c11 = {0};
  360. D c12 = {0};
  361. D c13 = {0};
  362. D c20 = {0};
  363. D c21 = {0};
  364. D c22 = {0};
  365. D c23 = {0};
  366. for (int l = 0; l < k; l += KN) {
  367. V k0 = load<V>(B + ldb * (j + 0) + l);
  368. V k1 = load<V>(B + ldb * (j + 1) + l);
  369. V k2 = load<V>(B + ldb * (j + 2) + l);
  370. V k3 = load<V>(B + ldb * (j + 3) + l);
  371. V a0 = load<V>(A + lda * (i + 0) + l);
  372. c00 = madd(a0, k0, c00);
  373. c01 = madd(a0, k1, c01);
  374. c02 = madd(a0, k2, c02);
  375. c03 = madd(a0, k3, c03);
  376. V a1 = load<V>(A + lda * (i + 1) + l);
  377. c10 = madd(a1, k0, c10);
  378. c11 = madd(a1, k1, c11);
  379. c12 = madd(a1, k2, c12);
  380. c13 = madd(a1, k3, c13);
  381. V a2 = load<V>(A + lda * (i + 2) + l);
  382. c20 = madd(a2, k0, c20);
  383. c21 = madd(a2, k1, c21);
  384. c22 = madd(a2, k2, c22);
  385. c23 = madd(a2, k3, c23);
  386. }
  387. C[ldc * (j + 0) + (i + 0)] = hsum(c00);
  388. C[ldc * (j + 0) + (i + 1)] = hsum(c10);
  389. C[ldc * (j + 0) + (i + 2)] = hsum(c20);
  390. C[ldc * (j + 1) + (i + 0)] = hsum(c01);
  391. C[ldc * (j + 1) + (i + 1)] = hsum(c11);
  392. C[ldc * (j + 1) + (i + 2)] = hsum(c21);
  393. C[ldc * (j + 2) + (i + 0)] = hsum(c02);
  394. C[ldc * (j + 2) + (i + 1)] = hsum(c12);
  395. C[ldc * (j + 2) + (i + 2)] = hsum(c22);
  396. C[ldc * (j + 3) + (i + 0)] = hsum(c03);
  397. C[ldc * (j + 3) + (i + 1)] = hsum(c13);
  398. C[ldc * (j + 3) + (i + 2)] = hsum(c23);
  399. END_KERNEL()
  400. }
  401. NOINLINE void gemm1x4(int m0, int m, int n0, int n) {
  402. BEGIN_KERNEL(1, 4)
  403. D c00 = {0}, e00 = {0};
  404. D c01 = {0}, e01 = {0};
  405. D c02 = {0}, e02 = {0};
  406. D c03 = {0}, e03 = {0};
  407. for (int l = 0; l < k; l += KN) {
  408. V a = load<V>(A + lda * (i + 0) + l);
  409. c00 = madder(a, load<V>(B + ldb * (j + 0) + l), c00, &e00);
  410. c01 = madder(a, load<V>(B + ldb * (j + 1) + l), c01, &e01);
  411. c02 = madder(a, load<V>(B + ldb * (j + 2) + l), c02, &e02);
  412. c03 = madder(a, load<V>(B + ldb * (j + 3) + l), c03, &e03);
  413. }
  414. C[ldc * (j + 0) + (i + 0)] = hsum(c00);
  415. C[ldc * (j + 1) + (i + 0)] = hsum(c01);
  416. C[ldc * (j + 2) + (i + 0)] = hsum(c02);
  417. C[ldc * (j + 3) + (i + 0)] = hsum(c03);
  418. END_KERNEL()
  419. }
  420. NOINLINE void gemm4x1(int m0, int m, int n0, int n) {
  421. BEGIN_KERNEL(4, 1)
  422. D c00 = {0}, e00 = {0};
  423. D c10 = {0}, e10 = {0};
  424. D c20 = {0}, e20 = {0};
  425. D c30 = {0}, e30 = {0};
  426. for (int l = 0; l < k; l += KN) {
  427. V b = load<V>(B + ldb * (j + 0) + l);
  428. c00 = madder(load<V>(A + lda * (i + 0) + l), b, c00, &e00);
  429. c10 = madder(load<V>(A + lda * (i + 1) + l), b, c10, &e10);
  430. c20 = madder(load<V>(A + lda * (i + 2) + l), b, c20, &e20);
  431. c30 = madder(load<V>(A + lda * (i + 3) + l), b, c30, &e30);
  432. }
  433. C[ldc * (j + 0) + (i + 0)] = hsum(c00);
  434. C[ldc * (j + 0) + (i + 1)] = hsum(c10);
  435. C[ldc * (j + 0) + (i + 2)] = hsum(c20);
  436. C[ldc * (j + 0) + (i + 3)] = hsum(c30);
  437. END_KERNEL()
  438. }
  439. NOINLINE void gemm1x1(int m0, int m, int n0, int n) {
  440. BEGIN_KERNEL(1, 1)
  441. D c = {0}, e = {0};
  442. for (int l = 0; l < k; l += KN)
  443. c = madder(load<V>(A + lda * i + l),
  444. load<V>(B + ldb * j + l), c, &e);
  445. C[ldc * j + i] = hsum(c);
  446. END_KERNEL()
  447. }
  448. const TA *const A;
  449. const TB *const B;
  450. TC *const C;
  451. const int k;
  452. const int lda;
  453. const int ldb;
  454. const int ldc;
  455. const int ith;
  456. const int nth;
  457. };
  458. //////////////////////////////////////////////////////////////////////////////////////////
  459. // QUANT ZERO MATRIX MULTIPLICATION
  460. #if defined(__ARM_FEATURE_DOTPROD)
  461. template <typename TA>
  462. class tinyBLAS_Q0_ARM {
  463. public:
  464. tinyBLAS_Q0_ARM(int k,
  465. const TA *A, int lda,
  466. const block_q8_0 *B, int ldb,
  467. float *C, int ldc,
  468. int ith, int nth)
  469. : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) {
  470. }
  471. void matmul(int m, int n, int task) {
  472. if (task == GGML_TASK_TYPE_COMPUTE)
  473. mnpack(0, m, 0, n);
  474. }
  475. private:
  476. NOINLINE void mnpack(int m0, int m, int n0, int n) {
  477. int mc, nc, mp, np;
  478. if (m - m0 <= 0 || n - n0 <= 0)
  479. return;
  480. if (m - m0 >= 3 && n - n0 >= 3) {
  481. mc = 3;
  482. nc = 3;
  483. gemm3x3(m0, m, n0, n);
  484. } else {
  485. mc = 1;
  486. nc = 1;
  487. gemm1x1(m0, m, n0, n);
  488. }
  489. mp = m0 + (m - m0) / mc * mc;
  490. np = n0 + (n - n0) / nc * nc;
  491. mnpack(mp, m, n0, np);
  492. mnpack(m0, mp, np, n);
  493. mnpack(mp, m, np, n);
  494. }
  495. NOINLINE void gemm3x3(int m0, int m, int n0, int n) {
  496. BEGIN_KERNEL(3, 3)
  497. int32x4_t zero = vdupq_n_s32(0);
  498. float32x4_t c00 = vdupq_n_f32(0.f);
  499. float32x4_t c01 = vdupq_n_f32(0.f);
  500. float32x4_t c02 = vdupq_n_f32(0.f);
  501. float32x4_t c10 = vdupq_n_f32(0.f);
  502. float32x4_t c11 = vdupq_n_f32(0.f);
  503. float32x4_t c12 = vdupq_n_f32(0.f);
  504. float32x4_t c20 = vdupq_n_f32(0.f);
  505. float32x4_t c21 = vdupq_n_f32(0.f);
  506. float32x4_t c22 = vdupq_n_f32(0.f);
  507. const TA *Ap0 = A + lda * (i + 0);
  508. const TA *Ap1 = A + lda * (i + 1);
  509. const TA *Ap2 = A + lda * (i + 2);
  510. const block_q8_0 *Bp0 = B + ldb * (j + 0);
  511. const block_q8_0 *Bp1 = B + ldb * (j + 1);
  512. const block_q8_0 *Bp2 = B + ldb * (j + 2);
  513. for (int l = 0; l < k; ++l) {
  514. c00 = vmlaq_n_f32(
  515. c00,
  516. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap0 + l), load_lo(Bp0 + l)),
  517. load_hi(Ap0 + l), load_hi(Bp0 + l))),
  518. unhalf(Ap0[l].d) * unhalf(Bp0[l].d));
  519. c01 = vmlaq_n_f32(
  520. c01,
  521. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap0 + l), load_lo(Bp1 + l)),
  522. load_hi(Ap0 + l), load_hi(Bp1 + l))),
  523. unhalf(Ap0[l].d) * unhalf(Bp1[l].d));
  524. c02 = vmlaq_n_f32(
  525. c02,
  526. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap0 + l), load_lo(Bp2 + l)),
  527. load_hi(Ap0 + l), load_hi(Bp2 + l))),
  528. unhalf(Ap0[l].d) * unhalf(Bp2[l].d));
  529. c10 = vmlaq_n_f32(
  530. c10,
  531. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap1 + l), load_lo(Bp0 + l)),
  532. load_hi(Ap1 + l), load_hi(Bp0 + l))),
  533. unhalf(Ap1[l].d) * unhalf(Bp0[l].d));
  534. c11 = vmlaq_n_f32(
  535. c11,
  536. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap1 + l), load_lo(Bp1 + l)),
  537. load_hi(Ap1 + l), load_hi(Bp1 + l))),
  538. unhalf(Ap1[l].d) * unhalf(Bp1[l].d));
  539. c12 = vmlaq_n_f32(
  540. c12,
  541. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap1 + l), load_lo(Bp2 + l)),
  542. load_hi(Ap1 + l), load_hi(Bp2 + l))),
  543. unhalf(Ap1[l].d) * unhalf(Bp2[l].d));
  544. c20 = vmlaq_n_f32(
  545. c20,
  546. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap2 + l), load_lo(Bp0 + l)),
  547. load_hi(Ap2 + l), load_hi(Bp0 + l))),
  548. unhalf(Ap2[l].d) * unhalf(Bp0[l].d));
  549. c21 = vmlaq_n_f32(
  550. c21,
  551. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap2 + l), load_lo(Bp1 + l)),
  552. load_hi(Ap2 + l), load_hi(Bp1 + l))),
  553. unhalf(Ap2[l].d) * unhalf(Bp1[l].d));
  554. c22 = vmlaq_n_f32(
  555. c22,
  556. vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap2 + l), load_lo(Bp2 + l)),
  557. load_hi(Ap2 + l), load_hi(Bp2 + l))),
  558. unhalf(Ap2[l].d) * unhalf(Bp2[l].d));
  559. }
  560. C[ldc * (j + 0) + (i + 0)] = hsum(c00);
  561. C[ldc * (j + 0) + (i + 1)] = hsum(c10);
  562. C[ldc * (j + 0) + (i + 2)] = hsum(c20);
  563. C[ldc * (j + 1) + (i + 0)] = hsum(c01);
  564. C[ldc * (j + 1) + (i + 1)] = hsum(c11);
  565. C[ldc * (j + 1) + (i + 2)] = hsum(c21);
  566. C[ldc * (j + 2) + (i + 0)] = hsum(c02);
  567. C[ldc * (j + 2) + (i + 1)] = hsum(c12);
  568. C[ldc * (j + 2) + (i + 2)] = hsum(c22);
  569. END_KERNEL()
  570. }
  571. NOINLINE void gemm1x1(int m0, int m, int n0, int n) {
  572. BEGIN_KERNEL(1, 1)
  573. float32x4_t acc = vdupq_n_f32(0.f);
  574. const TA *Ap = A + lda * i;
  575. const block_q8_0 *Bp = B + ldb * j;
  576. for (int l = 0; l < k; ++l) {
  577. acc = vmlaq_n_f32(acc,
  578. vcvtq_f32_s32(vdotq_s32(
  579. vdotq_s32(vdupq_n_s32(0), load_lo(Ap + l), load_lo(Bp + l)),
  580. load_hi(Ap + l), load_hi(Bp + l))),
  581. unhalf(Ap[l].d) * unhalf(Bp[l].d));
  582. }
  583. C[ldc * j + i] = hsum(acc);
  584. END_KERNEL()
  585. }
  586. inline int8x16_t load_lo(const block_q8_0 *b) {
  587. return vld1q_s8(b->qs);
  588. }
  589. inline int8x16_t load_hi(const block_q8_0 *b) {
  590. return vld1q_s8(b->qs + 16);
  591. }
  592. inline int8x16_t load_lo(const block_q4_0 *b) {
  593. return vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vld1q_u8(b->qs),
  594. vdupq_n_u8(0x0f))),
  595. vdupq_n_s8(0x8));
  596. }
  597. inline int8x16_t load_hi(const block_q4_0 *b) {
  598. return vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(vld1q_u8(b->qs), 4)),
  599. vdupq_n_s8(0x8));
  600. }
  601. const TA *const A;
  602. const block_q8_0 *const B;
  603. float *const C;
  604. const int k;
  605. const int lda;
  606. const int ldb;
  607. const int ldc;
  608. const int ith;
  609. const int nth;
  610. };
  611. #endif // __ARM_FEATURE_DOTPROD
  612. #if defined(__AVX2__) || defined(__AVX512F__)
  613. template <typename TA, typename TB, typename TC>
  614. class tinyBLAS_Q0_AVX2 {
  615. public:
  616. tinyBLAS_Q0_AVX2(int k,
  617. const TA *A, int lda,
  618. const TB *B, int ldb,
  619. TC *C, int ldc,
  620. int ith, int nth)
  621. : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) {
  622. }
  623. void matmul(int m, int n, int task) {
  624. if (task == GGML_TASK_TYPE_COMPUTE)
  625. mnpack(0, m, 0, n);
  626. }
  627. private:
  628. NOINLINE void mnpack(int m0, int m, int n0, int n) {
  629. int mc, nc, mp, np;
  630. if (m - m0 <= 0 || n - n0 <= 0)
  631. return;
  632. if (m - m0 >= 4 && n - n0 >= 3) {
  633. mc = 4;
  634. nc = 3;
  635. gemm4x3(m0, m, n0, n);
  636. } else if (m - m0 >= 4 && n - n0 >= 1) {
  637. mc = 4;
  638. nc = 1;
  639. gemm4x1(m0, m, n0, n);
  640. } else if (m - m0 >= 1 && n - n0 >= 4) {
  641. mc = 1;
  642. nc = 4;
  643. gemm1x4(m0, m, n0, n);
  644. } else {
  645. mc = 1;
  646. nc = 1;
  647. gemm1x1(m0, m, n0, n);
  648. }
  649. mp = m0 + (m - m0) / mc * mc;
  650. np = n0 + (n - n0) / nc * nc;
  651. mnpack(mp, m, n0, np);
  652. mnpack(m0, mp, np, n);
  653. mnpack(mp, m, np, n);
  654. }
  655. NOINLINE void gemm4x3(int m0, int m, int n0, int n) {
  656. BEGIN_KERNEL(4, 3)
  657. __m256 c00 = _mm256_setzero_ps();
  658. __m256 c10 = _mm256_setzero_ps();
  659. __m256 c20 = _mm256_setzero_ps();
  660. __m256 c30 = _mm256_setzero_ps();
  661. __m256 c01 = _mm256_setzero_ps();
  662. __m256 c11 = _mm256_setzero_ps();
  663. __m256 c21 = _mm256_setzero_ps();
  664. __m256 c31 = _mm256_setzero_ps();
  665. __m256 c02 = _mm256_setzero_ps();
  666. __m256 c12 = _mm256_setzero_ps();
  667. __m256 c22 = _mm256_setzero_ps();
  668. __m256 c32 = _mm256_setzero_ps();
  669. const TA *Ap0 = A + lda * (i + 0);
  670. const TA *Ap1 = A + lda * (i + 1);
  671. const TA *Ap2 = A + lda * (i + 2);
  672. const TA *Ap3 = A + lda * (i + 3);
  673. const TB *Bp0 = B + ldb * (j + 0);
  674. const TB *Bp1 = B + ldb * (j + 1);
  675. const TB *Bp2 = B + ldb * (j + 2);
  676. for (int l = 0; l < k; ++l) {
  677. float da0 = unhalf(Ap0[l].d);
  678. float da1 = unhalf(Ap1[l].d);
  679. float da2 = unhalf(Ap2[l].d);
  680. float da3 = unhalf(Ap3[l].d);
  681. __m256i e0 = load(Ap0 + l);
  682. __m256i e1 = load(Ap1 + l);
  683. __m256i e2 = load(Ap2 + l);
  684. __m256i e3 = load(Ap3 + l);
  685. float db0 = unhalf(Bp0[l].d);
  686. __m256 d00 = _mm256_set1_ps(da0 * db0);
  687. __m256 d10 = _mm256_set1_ps(da1 * db0);
  688. __m256 d20 = _mm256_set1_ps(da2 * db0);
  689. __m256 d30 = _mm256_set1_ps(da3 * db0);
  690. __m256i f0 = load(Bp0 + l);
  691. __m256i u0 = _mm256_sign_epi8(f0, f0);
  692. __m256i s00 = _mm256_sign_epi8(e0, f0);
  693. __m256i s10 = _mm256_sign_epi8(e1, f0);
  694. __m256i s20 = _mm256_sign_epi8(e2, f0);
  695. __m256i s30 = _mm256_sign_epi8(e3, f0);
  696. c00 = madd(d00, updot(u0, s00), c00);
  697. c10 = madd(d10, updot(u0, s10), c10);
  698. c20 = madd(d20, updot(u0, s20), c20);
  699. c30 = madd(d30, updot(u0, s30), c30);
  700. float db1 = unhalf(Bp1[l].d);
  701. __m256 d01 = _mm256_set1_ps(da0 * db1);
  702. __m256 d11 = _mm256_set1_ps(da1 * db1);
  703. __m256 d21 = _mm256_set1_ps(da2 * db1);
  704. __m256 d31 = _mm256_set1_ps(da3 * db1);
  705. __m256i f1 = load(Bp1 + l);
  706. __m256i u1 = _mm256_sign_epi8(f1, f1);
  707. __m256i s01 = _mm256_sign_epi8(e0, f1);
  708. __m256i s11 = _mm256_sign_epi8(e1, f1);
  709. __m256i s21 = _mm256_sign_epi8(e2, f1);
  710. __m256i s31 = _mm256_sign_epi8(e3, f1);
  711. c01 = madd(d01, updot(u1, s01), c01);
  712. c11 = madd(d11, updot(u1, s11), c11);
  713. c21 = madd(d21, updot(u1, s21), c21);
  714. c31 = madd(d31, updot(u1, s31), c31);
  715. float db2 = unhalf(Bp2[l].d);
  716. __m256 d02 = _mm256_set1_ps(da0 * db2);
  717. __m256 d12 = _mm256_set1_ps(da1 * db2);
  718. __m256 d22 = _mm256_set1_ps(da2 * db2);
  719. __m256 d32 = _mm256_set1_ps(da3 * db2);
  720. __m256i f2 = load(Bp2 + l);
  721. __m256i u2 = _mm256_sign_epi8(f2, f2);
  722. __m256i s02 = _mm256_sign_epi8(e0, f2);
  723. __m256i s12 = _mm256_sign_epi8(e1, f2);
  724. __m256i s22 = _mm256_sign_epi8(e2, f2);
  725. __m256i s32 = _mm256_sign_epi8(e3, f2);
  726. c02 = madd(d02, updot(u2, s02), c02);
  727. c12 = madd(d12, updot(u2, s12), c12);
  728. c22 = madd(d22, updot(u2, s22), c22);
  729. c32 = madd(d32, updot(u2, s32), c32);
  730. }
  731. C[ldc * (j + 0) + (i + 0)] = hsum(c00);
  732. C[ldc * (j + 0) + (i + 1)] = hsum(c10);
  733. C[ldc * (j + 0) + (i + 2)] = hsum(c20);
  734. C[ldc * (j + 0) + (i + 3)] = hsum(c30);
  735. C[ldc * (j + 1) + (i + 0)] = hsum(c01);
  736. C[ldc * (j + 1) + (i + 1)] = hsum(c11);
  737. C[ldc * (j + 1) + (i + 2)] = hsum(c21);
  738. C[ldc * (j + 1) + (i + 3)] = hsum(c31);
  739. C[ldc * (j + 2) + (i + 0)] = hsum(c02);
  740. C[ldc * (j + 2) + (i + 1)] = hsum(c12);
  741. C[ldc * (j + 2) + (i + 2)] = hsum(c22);
  742. C[ldc * (j + 2) + (i + 3)] = hsum(c32);
  743. END_KERNEL()
  744. }
  745. NOINLINE void gemm4x1(int m0, int m, int n0, int n) {
  746. BEGIN_KERNEL(4, 1)
  747. __m256 c0 = _mm256_setzero_ps();
  748. __m256 c1 = _mm256_setzero_ps();
  749. __m256 c2 = _mm256_setzero_ps();
  750. __m256 c3 = _mm256_setzero_ps();
  751. const TA *Ap0 = A + lda * (i + 0);
  752. const TA *Ap1 = A + lda * (i + 1);
  753. const TA *Ap2 = A + lda * (i + 2);
  754. const TA *Ap3 = A + lda * (i + 3);
  755. const TB *Bp = B + ldb * j;
  756. for (int l = 0; l < k; ++l) {
  757. float db0 = unhalf(Bp[l].d);
  758. __m256i f = load(Bp + l);
  759. __m256i u = _mm256_sign_epi8(f, f);
  760. __m256 d0 = _mm256_set1_ps(unhalf(Ap0[l].d) * db0);
  761. __m256 d1 = _mm256_set1_ps(unhalf(Ap1[l].d) * db0);
  762. __m256 d2 = _mm256_set1_ps(unhalf(Ap2[l].d) * db0);
  763. __m256 d3 = _mm256_set1_ps(unhalf(Ap3[l].d) * db0);
  764. __m256i e0 = load(Ap0 + l);
  765. __m256i e1 = load(Ap1 + l);
  766. __m256i e2 = load(Ap2 + l);
  767. __m256i e3 = load(Ap3 + l);
  768. __m256i s0 = _mm256_sign_epi8(e0, f);
  769. __m256i s1 = _mm256_sign_epi8(e1, f);
  770. __m256i s2 = _mm256_sign_epi8(e2, f);
  771. __m256i s3 = _mm256_sign_epi8(e3, f);
  772. __m256 g0 = updot(u, s0);
  773. __m256 g1 = updot(u, s1);
  774. __m256 g2 = updot(u, s2);
  775. __m256 g3 = updot(u, s3);
  776. c0 = madd(d0, g0, c0);
  777. c1 = madd(d1, g1, c1);
  778. c2 = madd(d2, g2, c2);
  779. c3 = madd(d3, g3, c3);
  780. }
  781. C[ldc * j + (i + 0)] = hsum(c0);
  782. C[ldc * j + (i + 1)] = hsum(c1);
  783. C[ldc * j + (i + 2)] = hsum(c2);
  784. C[ldc * j + (i + 3)] = hsum(c3);
  785. END_KERNEL()
  786. }
  787. NOINLINE void gemm1x4(int m0, int m, int n0, int n) {
  788. BEGIN_KERNEL(1, 4)
  789. __m256 c0 = _mm256_setzero_ps();
  790. __m256 c1 = _mm256_setzero_ps();
  791. __m256 c2 = _mm256_setzero_ps();
  792. __m256 c3 = _mm256_setzero_ps();
  793. const TB *Bp0 = B + ldb * (j + 0);
  794. const TB *Bp1 = B + ldb * (j + 1);
  795. const TB *Bp2 = B + ldb * (j + 2);
  796. const TB *Bp3 = B + ldb * (j + 3);
  797. const TA *Ap = A + lda * i;
  798. for (int l = 0; l < k; ++l) {
  799. float da0 = unhalf(Ap[l].d);
  800. __m256i f = load(Ap + l);
  801. __m256i u = _mm256_sign_epi8(f, f);
  802. __m256 d0 = _mm256_set1_ps(unhalf(Bp0[l].d) * da0);
  803. __m256 d1 = _mm256_set1_ps(unhalf(Bp1[l].d) * da0);
  804. __m256 d2 = _mm256_set1_ps(unhalf(Bp2[l].d) * da0);
  805. __m256 d3 = _mm256_set1_ps(unhalf(Bp3[l].d) * da0);
  806. __m256 g0 = updot(u, _mm256_sign_epi8(load(Bp0 + l), f));
  807. __m256 g1 = updot(u, _mm256_sign_epi8(load(Bp1 + l), f));
  808. __m256 g2 = updot(u, _mm256_sign_epi8(load(Bp2 + l), f));
  809. __m256 g3 = updot(u, _mm256_sign_epi8(load(Bp3 + l), f));
  810. c0 = madd(d0, g0, c0);
  811. c1 = madd(d1, g1, c1);
  812. c2 = madd(d2, g2, c2);
  813. c3 = madd(d3, g3, c3);
  814. }
  815. C[ldc * (j + 0) + i] = hsum(c0);
  816. C[ldc * (j + 1) + i] = hsum(c1);
  817. C[ldc * (j + 2) + i] = hsum(c2);
  818. C[ldc * (j + 3) + i] = hsum(c3);
  819. END_KERNEL()
  820. }
  821. NOINLINE void gemm1x1(int m0, int m, int n0, int n) {
  822. BEGIN_KERNEL(1, 1)
  823. __m256 c = _mm256_setzero_ps();
  824. const TA *Ap = A + lda * i;
  825. const TB *Bp = B + ldb * j;
  826. for (int l = 0; l < k; ++l) {
  827. __m256 d = _mm256_set1_ps(unhalf(Ap[l].d) * unhalf(Bp[l].d));
  828. __m256i e = load(Ap + l);
  829. __m256i f = load(Bp + l);
  830. __m256 g = updot(_mm256_sign_epi8(e, e), _mm256_sign_epi8(f, e));
  831. c = madd(d, g, c);
  832. }
  833. C[ldc * j + i] = hsum(c);
  834. END_KERNEL()
  835. }
  836. inline __m256i load(const block_q8_0 *b) {
  837. return _mm256_loadu_si256((const __m256i *)b->qs);
  838. }
  839. inline __m256i load(const block_q4_0 *b) {
  840. return _mm256_sub_epi8(denibble(b->qs), _mm256_set1_epi8(8));
  841. }
  842. inline __m256 updot(__m256i u, __m256i s) {
  843. __m256i res;
  844. #if defined(__AVXVNNI__) || (defined(__AVX512VNNI__) && defined(__AVX512VL__))
  845. res = _mm256_dpbusd_epi32(_mm256_setzero_si256(), u, s);
  846. #else
  847. res = _mm256_madd_epi16(_mm256_set1_epi16(1), _mm256_maddubs_epi16(u, s));
  848. #endif
  849. return _mm256_cvtepi32_ps(res);
  850. }
  851. static inline __m256i denibble(const uint8_t *p) {
  852. const __m128i tmp = _mm_loadu_si128((const __m128i *)p);
  853. const __m256i bytes = MM256_SET_M128I(_mm_srli_epi16(tmp, 4), tmp);
  854. const __m256i lowMask = _mm256_set1_epi8(15);
  855. return _mm256_and_si256(lowMask, bytes);
  856. }
  857. const TA *const A;
  858. const TB *const B;
  859. TC *const C;
  860. const int k;
  861. const int lda;
  862. const int ldb;
  863. const int ldc;
  864. const int ith;
  865. const int nth;
  866. };
  867. #endif // __AVX2__
  868. } // namespace
  869. /**
  870. * Performs optimized matrix multiplication on CPU.
  871. *
  872. * This subroutine may compute C = Aᵀ * B with column major ordering.
  873. * Despite its name, this isn't a generalized implementation. Work is
  874. * only performed when a handwritten kernel is written and available.
  875. * Otherwise the caller should fall back to a general matmul routine.
  876. *
  877. * For example, for single-threaded single-precision GEMM you can say
  878. *
  879. * llamafile_sgemm(m, n, k, A, lda, B, ldb, C, ldc,
  880. * 0, 1, GGML_TASK_TYPE_COMPUTE,
  881. * GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32);
  882. *
  883. * @param m is rows in `A` and `C`
  884. * @param n is cols in `B` and `C`
  885. * @param k is cols in `A` and rows in `B`
  886. * @param A is first input matrix (always transposed)
  887. * @param lda is row stride of `A`
  888. * @param B is second input matrix (never transposed)
  889. * @param ldb is row stride of `B`
  890. * @param C is input/output array of output matrices
  891. * @param ldc is row stride of `C`
  892. * @param ith is thread id (must be less than `nth`)
  893. * @param nth is number of threads (must be greater than zero)
  894. * @param task is GGML task type
  895. * @param Atype is GGML data type of `A`
  896. * @param Btype is GGML data type of `B`
  897. * @param Ctype is GGML data type of `C`
  898. * @return true if this function was able to service the matmul request
  899. */
  900. bool llamafile_sgemm(int m, int n, int k, const void *A, int lda, const void *B, int ldb, void *C,
  901. int ldc, int ith, int nth, int task, int Atype, int Btype, int Ctype) {
  902. assert(m >= 0);
  903. assert(n >= 0);
  904. assert(k >= 0);
  905. assert(lda >= k);
  906. assert(ldb >= k);
  907. assert(ldc >= m);
  908. assert(nth > 0);
  909. assert(ith < nth);
  910. assert(1ll * lda * m <= 0x7fffffff);
  911. assert(1ll * ldb * n <= 0x7fffffff);
  912. assert(1ll * ldc * n <= 0x7fffffff);
  913. if (Ctype != GGML_TYPE_F32)
  914. return false;
  915. switch (Atype) {
  916. case GGML_TYPE_F32: {
  917. if (Btype != GGML_TYPE_F32)
  918. return false;
  919. #if defined(__AVX512F__)
  920. if (k % 16)
  921. return false;
  922. tinyBLAS<16, __m512, __m512, float, float, float> tb{
  923. k, (const float *)A, lda,
  924. (const float *)B, ldb,
  925. (float *)C, ldc,
  926. ith, nth};
  927. tb.matmul(m, n, task);
  928. return true;
  929. #elif defined(__AVX__) || defined(__AVX2__)
  930. if (k % 8)
  931. return false;
  932. tinyBLAS<8, __m256, __m256, float, float, float> tb{
  933. k, (const float *)A, lda,
  934. (const float *)B, ldb,
  935. (float *)C, ldc,
  936. ith, nth};
  937. tb.matmul(m, n, task);
  938. return true;
  939. #elif defined(__ARM_NEON)
  940. if (n < 4)
  941. return false;
  942. if (k % 4)
  943. return false;
  944. tinyBLAS<4, float32x4_t, float32x4_t, float, float, float> tb{
  945. k, (const float *)A, lda,
  946. (const float *)B, ldb,
  947. (float *)C, ldc,
  948. ith, nth};
  949. tb.matmul(m, n, task);
  950. return true;
  951. #else
  952. return false;
  953. #endif
  954. }
  955. case GGML_TYPE_F16: {
  956. #if defined(__AVX512F__)
  957. if (k % 16)
  958. return false;
  959. if (Btype != GGML_TYPE_F32)
  960. return false;
  961. tinyBLAS<16, __m512, __m512, ggml_fp16_t, float, float> tb{
  962. k, (const ggml_fp16_t *)A, lda,
  963. (const float *)B, ldb,
  964. (float *)C, ldc,
  965. ith, nth};
  966. tb.matmul(m, n, task);
  967. return true;
  968. #elif (defined(__AVX__) || defined(__AVX2__)) && defined(__F16C__)
  969. if (k % 8)
  970. return false;
  971. if (Btype != GGML_TYPE_F32)
  972. return false;
  973. tinyBLAS<8, __m256, __m256, ggml_fp16_t, float, float> tb{
  974. k, (const ggml_fp16_t *)A, lda,
  975. (const float *)B, ldb,
  976. (float *)C, ldc,
  977. ith, nth};
  978. tb.matmul(m, n, task);
  979. return true;
  980. #elif defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER)
  981. if (n < 8)
  982. return false;
  983. if (k % 8)
  984. return false;
  985. if (Btype != GGML_TYPE_F16)
  986. return false;
  987. tinyBLAS<8, float16x8_t, float16x8_t, ggml_fp16_t, ggml_fp16_t, float> tb{
  988. k, (const ggml_fp16_t *)A, lda,
  989. (const ggml_fp16_t *)B, ldb,
  990. (float *)C, ldc,
  991. ith, nth};
  992. tb.matmul(m, n, task);
  993. return true;
  994. #elif defined(__ARM_NEON) && !defined(_MSC_VER)
  995. if (k % 4)
  996. return false;
  997. if (Btype != GGML_TYPE_F32)
  998. return false;
  999. tinyBLAS<4, float32x4_t, float32x4_t, ggml_fp16_t, float, float> tb{
  1000. k, (const ggml_fp16_t *)A, lda,
  1001. (const float *)B, ldb,
  1002. (float *)C, ldc,
  1003. ith, nth};
  1004. tb.matmul(m, n, task);
  1005. return true;
  1006. #else
  1007. return false;
  1008. #endif
  1009. }
  1010. case GGML_TYPE_Q8_0: {
  1011. if (Btype != GGML_TYPE_Q8_0)
  1012. return false;
  1013. #if defined(__AVX2__) || defined(__AVX512F__)
  1014. tinyBLAS_Q0_AVX2<block_q8_0, block_q8_0, float> tb{
  1015. k, (const block_q8_0 *)A, lda,
  1016. (const block_q8_0 *)B, ldb,
  1017. (float *)C, ldc,
  1018. ith, nth};
  1019. tb.matmul(m, n, task);
  1020. return true;
  1021. #elif defined(__ARM_FEATURE_DOTPROD)
  1022. tinyBLAS_Q0_ARM<block_q8_0> tb{
  1023. k, (const block_q8_0 *)A, lda,
  1024. (const block_q8_0 *)B, ldb,
  1025. (float *)C, ldc,
  1026. ith, nth};
  1027. tb.matmul(m, n, task);
  1028. return true;
  1029. #else
  1030. return false;
  1031. #endif
  1032. }
  1033. case GGML_TYPE_Q4_0: {
  1034. if (Btype != GGML_TYPE_Q8_0)
  1035. return false;
  1036. #if defined(__AVX2__) || defined(__AVX512F__)
  1037. tinyBLAS_Q0_AVX2<block_q4_0, block_q8_0, float> tb{
  1038. k, (const block_q4_0 *)A, lda,
  1039. (const block_q8_0 *)B, ldb,
  1040. (float *)C, ldc,
  1041. ith, nth};
  1042. tb.matmul(m, n, task);
  1043. return true;
  1044. #elif defined(__ARM_FEATURE_DOTPROD)
  1045. tinyBLAS_Q0_ARM<block_q4_0> tb{
  1046. k, (const block_q4_0 *)A, lda,
  1047. (const block_q8_0 *)B, ldb,
  1048. (float *)C, ldc,
  1049. ith, nth};
  1050. tb.matmul(m, n, task);
  1051. return true;
  1052. #else
  1053. return false;
  1054. #endif
  1055. }
  1056. default:
  1057. return false;
  1058. }
  1059. (void)m;
  1060. (void)n;
  1061. (void)k;
  1062. (void)A;
  1063. (void)lda;
  1064. (void)B;
  1065. (void)ldb;
  1066. (void)C;
  1067. (void)ldc;
  1068. (void)ith;
  1069. (void)nth;
  1070. (void)task;
  1071. (void)Atype;
  1072. (void)Btype;
  1073. (void)Ctype;
  1074. }