diff options
Diffstat (limited to 'sgemm.cpp')
| -rw-r--r-- | sgemm.cpp | 1148 |
1 files changed, 1148 insertions, 0 deletions
diff --git a/sgemm.cpp b/sgemm.cpp new file mode 100644 index 00000000..6900f04c --- /dev/null +++ b/sgemm.cpp @@ -0,0 +1,1148 @@ +// -*- mode:c++;indent-tabs-mode:nil;c-basic-offset:4;coding:utf-8 -*- +// vi: set et ft=c++ ts=4 sts=4 sw=4 fenc=utf-8 :vi +// +// Copyright 2024 Mozilla Foundation +// +// Permission is hereby granted, free of charge, to any person obtaining +// a copy of this software and associated documentation files (the +// "Software"), to deal in the Software without restriction, including +// without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to +// permit persons to whom the Software is furnished to do so, subject to +// the following conditions: +// +// The above copyright notice and this permission notice shall be +// included in all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS +// BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN +// ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +// SOFTWARE. + +// +// _ _ ___ _ _ ___ +// | |_(_)_ _ _ _| _ ) | /_\ / __| +// | _| | ' \ || | _ \ |__ / _ \\__ \. +// \__|_|_||_\_, |___/____/_/ \_\___/ +// |__/ +// +// BASIC LINEAR ALGEBRA SUBPROGRAMS +// +// +// This file implements multithreaded CPU matrix multiplication for the +// common contiguous use case C = Aᵀ * B. These kernels are designed to +// have excellent performance[1] for matrices that fit in the CPU cache +// without imposing any overhead such as cache filling or malloc calls. +// +// This implementation does not guarantee any upper bound with rounding +// errors, which grow along with k. Our goal's to maximally exploit the +// hardware for performance, and then use whatever resources remain for +// improving numerical accuracy. +// +// [1] J. Tunney, ‘LLaMA Now Goes Faster on CPUs’, Mar. 2024. [Online]. +// Available: https://justine.lol/matmul/. [Accessed: 29-Mar-2024]. + +#pragma GCC diagnostic ignored "-Wpedantic" +#pragma GCC diagnostic ignored "-Wignored-attributes" + +#include "sgemm.h" +#include "ggml-impl.h" +#include "ggml-quants.h" + +#ifdef _MSC_VER +#define NOINLINE __declspec(noinline) +#else +#define NOINLINE __attribute__((__noinline__)) +#endif + +#if defined(__ARM_NEON) || defined(__AVX512F__) +#define VECTOR_REGISTERS 32 +#else +#define VECTOR_REGISTERS 16 +#endif + +// there will be blocks +#define BEGIN_KERNEL(RM, RN) \ + int ytiles = (m - m0) / RM; \ + int xtiles = (n - n0) / RN; \ + int tiles = ytiles * xtiles; \ + int duty = (tiles + nth - 1) / nth; \ + int start = duty * ith; \ + int end = start + duty; \ + if (end > tiles) \ + end = tiles; \ + for (int job = start; job < end; ++job) { \ + int i = m0 + job / xtiles * RM; \ + int j = n0 + job % xtiles * RN; + +#define END_KERNEL() } + +#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) + +namespace { + +inline float unhalf(ggml_fp16_t d) { + return GGML_FP16_TO_FP32(d); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// VECTORIZED ARITHMETIC OPERATIONS + +#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +inline __m128 add(__m128 x, __m128 y) { return _mm_add_ps(x, y); } +inline __m128 sub(__m128 x, __m128 y) { return _mm_sub_ps(x, y); } +inline __m128 mul(__m128 x, __m128 y) { return _mm_mul_ps(x, y); } +#endif // __SSE__ + +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +inline __m256 add(__m256 x, __m256 y) { return _mm256_add_ps(x, y); } +inline __m256 sub(__m256 x, __m256 y) { return _mm256_sub_ps(x, y); } +inline __m256 mul(__m256 x, __m256 y) { return _mm256_mul_ps(x, y); } +#endif // __AVX__ + +#if defined(__AVX512F__) +inline __m512 add(__m512 x, __m512 y) { return _mm512_add_ps(x, y); } +inline __m512 sub(__m512 x, __m512 y) { return _mm512_sub_ps(x, y); } +inline __m512 mul(__m512 x, __m512 y) { return _mm512_mul_ps(x, y); } +#endif // __AVX512F__ + +#if defined(__ARM_NEON) +inline float32x4_t add(float32x4_t x, float32x4_t y) { return vaddq_f32(x, y); } +inline float32x4_t sub(float32x4_t x, float32x4_t y) { return vsubq_f32(x, y); } +inline float32x4_t mul(float32x4_t x, float32x4_t y) { return vmulq_f32(x, y); } +#endif // __ARM_NEON + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) +inline float16x8_t add(float16x8_t x, float16x8_t y) { return vaddq_f16(x, y); } +inline float16x8_t sub(float16x8_t x, float16x8_t y) { return vsubq_f16(x, y); } +inline float16x8_t mul(float16x8_t x, float16x8_t y) { return vmulq_f16(x, y); } +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// VECTORIZED HORIZONTAL SUM + +#if defined(__ARM_NEON) +inline float hsum(float32x4_t x) { + return vaddvq_f32(x); +} +#endif // __ARM_NEON + +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) +inline float hsum(float16x8_t x) { + return vaddvq_f32(vaddq_f32(vcvt_f32_f16(vget_low_f16(x)), + vcvt_f32_f16(vget_high_f16(x)))); +} +#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC + +#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +inline float hsum(__m128 x) { +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) + x = _mm_add_ps(x, _mm_movehl_ps(x, x)); + x = _mm_add_ss(x, _mm_movehdup_ps(x)); +#else + __m128 t; + t = _mm_shuffle_ps(x, x, _MM_SHUFFLE(2, 3, 0, 1)); + x = _mm_add_ps(x, t); + t = _mm_movehl_ps(t, x); + x = _mm_add_ss(x, t); +#endif + return _mm_cvtss_f32(x); +} +#endif + +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +inline float hsum(__m256 x) { + return hsum(_mm_add_ps(_mm256_extractf128_ps(x, 1), + _mm256_castps256_ps128(x))); +} +#endif // __AVX__ + +#if defined(__AVX512F__) +inline float hsum(__m512 x) { + return _mm512_reduce_add_ps(x); +} +#endif // __AVX512F__ + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// VECTORIZED MEMORY LOADING + +template <typename T, typename U> T load(const U *); + +#if defined(__ARM_NEON) +template <> inline float32x4_t load(const float *p) { + return vld1q_f32(p); +} +#if !defined(_MSC_VER) +template <> inline float16x8_t load(const ggml_fp16_t *p) { + return vld1q_f16((const float16_t *)p); +} +template <> inline float32x4_t load(const ggml_fp16_t *p) { + return vcvt_f32_f16(vld1_f16((const float16_t *)p)); +} +#endif // _MSC_VER +#endif // __ARM_NEON + +#if defined(__SSE__) || defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +template <> inline __m128 load(const float *p) { + return _mm_loadu_ps(p); +} +#endif // __SSE__ + +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) +template <> inline __m256 load(const float *p) { + return _mm256_loadu_ps(p); +} +#endif // __AVX__ + +#if defined(__F16C__) +template <> inline __m256 load(const ggml_fp16_t *p) { + return _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)p)); +} +#endif // __F16C__ + +#if defined(__AVX512F__) +template <> inline __m512 load(const float *p) { + return _mm512_loadu_ps(p); +} +template <> inline __m512 load(const ggml_fp16_t *p) { + return _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)p)); +} +#endif // __AVX512F__ + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// ABSTRACTIONS + +/** + * Computes a * b + c. + * + * This operation will become fused into a single arithmetic instruction + * if the hardware has support for this feature, e.g. Intel Haswell+ (c. + * 2013), AMD Bulldozer+ (c. 2011), etc. + */ +template <typename T, typename U> +inline U madd(T a, T b, U c) { + return add(mul(a, b), c); +} + +/** + * Computes a * b + c with error correction. + * + * @see W. Kahan, "Further remarks on reducing truncation errors," + * Communications of the ACM, vol. 8, no. 1, p. 40, Jan. 1965, + * doi: 10.1145/363707.363723. + */ +template <typename T, typename U> +inline U madder(T a, T b, U c, U *e) { + U y = sub(mul(a, b), *e); + U t = add(c, y); + *e = sub(sub(t, c), y); + return t; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// FLOATING POINT MATRIX MULTIPLICATION + +template <int KN, typename D, typename V, typename TA, typename TB, typename TC> +class tinyBLAS { + public: + tinyBLAS(int k, + const TA *A, int lda, + const TB *B, int ldb, + TC *C, int ldc, + int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { + } + + void matmul(int m, int n, int task) { + if (task == GGML_TASK_TYPE_COMPUTE) + mnpack(0, m, 0, n); + } + + private: + NOINLINE void mnpack(int m0, int m, int n0, int n) { + int mc, nc, mp, np; + if (m - m0 <= 0 || n - n0 <= 0) + return; + if (VECTOR_REGISTERS >= 32 && n - n0 >= 5 && m - m0 >= 5) { + mc = 5; + nc = 5; + gemm5x5(m0, m, n0, n); + } else if (n - n0 >= 4 && m - m0 >= 3) { + mc = 3; + nc = 4; + gemm3x4(m0, m, n0, n); + } else if (n - n0 >= 4) { + mc = 1; + nc = 4; + gemm1x4(m0, m, n0, n); + } else if (m - m0 >= 4) { + mc = 4; + nc = 1; + gemm4x1(m0, m, n0, n); + } else { + mc = 1; + nc = 1; + gemm1x1(m0, m, n0, n); + } + mp = m0 + (m - m0) / mc * mc; + np = n0 + (n - n0) / nc * nc; + mnpack(mp, m, n0, np); + mnpack(m0, mp, np, n); + mnpack(mp, m, np, n); + } + + NOINLINE void gemm5x5(int m0, int m, int n0, int n) { + BEGIN_KERNEL(5, 5) + D c00 = {0}; + D c01 = {0}; + D c02 = {0}; + D c03 = {0}; + D c04 = {0}; + D c10 = {0}; + D c11 = {0}; + D c12 = {0}; + D c13 = {0}; + D c14 = {0}; + D c20 = {0}; + D c21 = {0}; + D c22 = {0}; + D c23 = {0}; + D c24 = {0}; + D c30 = {0}; + D c31 = {0}; + D c32 = {0}; + D c33 = {0}; + D c34 = {0}; + D c40 = {0}; + D c41 = {0}; + D c42 = {0}; + D c43 = {0}; + D c44 = {0}; + for (int l = 0; l < k; l += KN) { + V k0 = load<V>(B + ldb * (j + 0) + l); + V k1 = load<V>(B + ldb * (j + 1) + l); + V k2 = load<V>(B + ldb * (j + 2) + l); + V k3 = load<V>(B + ldb * (j + 3) + l); + V k4 = load<V>(B + ldb * (j + 4) + l); + V a0 = load<V>(A + lda * (i + 0) + l); + c00 = madd(a0, k0, c00); + c01 = madd(a0, k1, c01); + c02 = madd(a0, k2, c02); + c03 = madd(a0, k3, c03); + c04 = madd(a0, k4, c04); + V a1 = load<V>(A + lda * (i + 1) + l); + c10 = madd(a1, k0, c10); + c11 = madd(a1, k1, c11); + c12 = madd(a1, k2, c12); + c13 = madd(a1, k3, c13); + c14 = madd(a1, k4, c14); + V a2 = load<V>(A + lda * (i + 2) + l); + c20 = madd(a2, k0, c20); + c21 = madd(a2, k1, c21); + c22 = madd(a2, k2, c22); + c23 = madd(a2, k3, c23); + c24 = madd(a2, k4, c24); + V a3 = load<V>(A + lda * (i + 3) + l); + c30 = madd(a3, k0, c30); + c31 = madd(a3, k1, c31); + c32 = madd(a3, k2, c32); + c33 = madd(a3, k3, c33); + c34 = madd(a3, k4, c34); + V a4 = load<V>(A + lda * (i + 4) + l); + c40 = madd(a4, k0, c40); + c41 = madd(a4, k1, c41); + c42 = madd(a4, k2, c42); + c43 = madd(a4, k3, c43); + c44 = madd(a4, k4, c44); + } + C[ldc * (j + 0) + (i + 0)] = hsum(c00); + C[ldc * (j + 0) + (i + 1)] = hsum(c10); + C[ldc * (j + 0) + (i + 2)] = hsum(c20); + C[ldc * (j + 0) + (i + 3)] = hsum(c30); + C[ldc * (j + 0) + (i + 4)] = hsum(c40); + C[ldc * (j + 1) + (i + 0)] = hsum(c01); + C[ldc * (j + 1) + (i + 1)] = hsum(c11); + C[ldc * (j + 1) + (i + 2)] = hsum(c21); + C[ldc * (j + 1) + (i + 3)] = hsum(c31); + C[ldc * (j + 1) + (i + 4)] = hsum(c41); + C[ldc * (j + 2) + (i + 0)] = hsum(c02); + C[ldc * (j + 2) + (i + 1)] = hsum(c12); + C[ldc * (j + 2) + (i + 2)] = hsum(c22); + C[ldc * (j + 2) + (i + 3)] = hsum(c32); + C[ldc * (j + 2) + (i + 4)] = hsum(c42); + C[ldc * (j + 3) + (i + 0)] = hsum(c03); + C[ldc * (j + 3) + (i + 1)] = hsum(c13); + C[ldc * (j + 3) + (i + 2)] = hsum(c23); + C[ldc * (j + 3) + (i + 3)] = hsum(c33); + C[ldc * (j + 3) + (i + 4)] = hsum(c43); + C[ldc * (j + 4) + (i + 0)] = hsum(c04); + C[ldc * (j + 4) + (i + 1)] = hsum(c14); + C[ldc * (j + 4) + (i + 2)] = hsum(c24); + C[ldc * (j + 4) + (i + 3)] = hsum(c34); + C[ldc * (j + 4) + (i + 4)] = hsum(c44); + END_KERNEL() + } + + NOINLINE void gemm3x4(int m0, int m, int n0, int n) { + BEGIN_KERNEL(3, 4) + D c00 = {0}; + D c01 = {0}; + D c02 = {0}; + D c03 = {0}; + D c10 = {0}; + D c11 = {0}; + D c12 = {0}; + D c13 = {0}; + D c20 = {0}; + D c21 = {0}; + D c22 = {0}; + D c23 = {0}; + for (int l = 0; l < k; l += KN) { + V k0 = load<V>(B + ldb * (j + 0) + l); + V k1 = load<V>(B + ldb * (j + 1) + l); + V k2 = load<V>(B + ldb * (j + 2) + l); + V k3 = load<V>(B + ldb * (j + 3) + l); + V a0 = load<V>(A + lda * (i + 0) + l); + c00 = madd(a0, k0, c00); + c01 = madd(a0, k1, c01); + c02 = madd(a0, k2, c02); + c03 = madd(a0, k3, c03); + V a1 = load<V>(A + lda * (i + 1) + l); + c10 = madd(a1, k0, c10); + c11 = madd(a1, k1, c11); + c12 = madd(a1, k2, c12); + c13 = madd(a1, k3, c13); + V a2 = load<V>(A + lda * (i + 2) + l); + c20 = madd(a2, k0, c20); + c21 = madd(a2, k1, c21); + c22 = madd(a2, k2, c22); + c23 = madd(a2, k3, c23); + } + C[ldc * (j + 0) + (i + 0)] = hsum(c00); + C[ldc * (j + 0) + (i + 1)] = hsum(c10); + C[ldc * (j + 0) + (i + 2)] = hsum(c20); + C[ldc * (j + 1) + (i + 0)] = hsum(c01); + C[ldc * (j + 1) + (i + 1)] = hsum(c11); + C[ldc * (j + 1) + (i + 2)] = hsum(c21); + C[ldc * (j + 2) + (i + 0)] = hsum(c02); + C[ldc * (j + 2) + (i + 1)] = hsum(c12); + C[ldc * (j + 2) + (i + 2)] = hsum(c22); + C[ldc * (j + 3) + (i + 0)] = hsum(c03); + C[ldc * (j + 3) + (i + 1)] = hsum(c13); + C[ldc * (j + 3) + (i + 2)] = hsum(c23); + END_KERNEL() + } + + NOINLINE void gemm1x4(int m0, int m, int n0, int n) { + BEGIN_KERNEL(1, 4) + D c00 = {0}, e00 = {0}; + D c01 = {0}, e01 = {0}; + D c02 = {0}, e02 = {0}; + D c03 = {0}, e03 = {0}; + for (int l = 0; l < k; l += KN) { + V a = load<V>(A + lda * (i + 0) + l); + c00 = madder(a, load<V>(B + ldb * (j + 0) + l), c00, &e00); + c01 = madder(a, load<V>(B + ldb * (j + 1) + l), c01, &e01); + c02 = madder(a, load<V>(B + ldb * (j + 2) + l), c02, &e02); + c03 = madder(a, load<V>(B + ldb * (j + 3) + l), c03, &e03); + } + C[ldc * (j + 0) + (i + 0)] = hsum(c00); + C[ldc * (j + 1) + (i + 0)] = hsum(c01); + C[ldc * (j + 2) + (i + 0)] = hsum(c02); + C[ldc * (j + 3) + (i + 0)] = hsum(c03); + END_KERNEL() + } + + NOINLINE void gemm4x1(int m0, int m, int n0, int n) { + BEGIN_KERNEL(4, 1) + D c00 = {0}, e00 = {0}; + D c10 = {0}, e10 = {0}; + D c20 = {0}, e20 = {0}; + D c30 = {0}, e30 = {0}; + for (int l = 0; l < k; l += KN) { + V b = load<V>(B + ldb * (j + 0) + l); + c00 = madder(load<V>(A + lda * (i + 0) + l), b, c00, &e00); + c10 = madder(load<V>(A + lda * (i + 1) + l), b, c10, &e10); + c20 = madder(load<V>(A + lda * (i + 2) + l), b, c20, &e20); + c30 = madder(load<V>(A + lda * (i + 3) + l), b, c30, &e30); + } + C[ldc * (j + 0) + (i + 0)] = hsum(c00); + C[ldc * (j + 0) + (i + 1)] = hsum(c10); + C[ldc * (j + 0) + (i + 2)] = hsum(c20); + C[ldc * (j + 0) + (i + 3)] = hsum(c30); + END_KERNEL() + } + + NOINLINE void gemm1x1(int m0, int m, int n0, int n) { + BEGIN_KERNEL(1, 1) + D c = {0}, e = {0}; + for (int l = 0; l < k; l += KN) + c = madder(load<V>(A + lda * i + l), + load<V>(B + ldb * j + l), c, &e); + C[ldc * j + i] = hsum(c); + END_KERNEL() + } + + const TA *const A; + const TB *const B; + TC *const C; + const int k; + const int lda; + const int ldb; + const int ldc; + const int ith; + const int nth; +}; + +////////////////////////////////////////////////////////////////////////////////////////// +// QUANT ZERO MATRIX MULTIPLICATION + +#if defined(__ARM_FEATURE_DOTPROD) +template <typename TA> +class tinyBLAS_Q0_ARM { + public: + tinyBLAS_Q0_ARM(int k, + const TA *A, int lda, + const block_q8_0 *B, int ldb, + float *C, int ldc, + int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { + } + + void matmul(int m, int n, int task) { + if (task == GGML_TASK_TYPE_COMPUTE) + mnpack(0, m, 0, n); + } + + private: + NOINLINE void mnpack(int m0, int m, int n0, int n) { + int mc, nc, mp, np; + if (m - m0 <= 0 || n - n0 <= 0) + return; + if (m - m0 >= 3 && n - n0 >= 3) { + mc = 3; + nc = 3; + gemm3x3(m0, m, n0, n); + } else { + mc = 1; + nc = 1; + gemm1x1(m0, m, n0, n); + } + mp = m0 + (m - m0) / mc * mc; + np = n0 + (n - n0) / nc * nc; + mnpack(mp, m, n0, np); + mnpack(m0, mp, np, n); + mnpack(mp, m, np, n); + } + + NOINLINE void gemm3x3(int m0, int m, int n0, int n) { + BEGIN_KERNEL(3, 3) + int32x4_t zero = vdupq_n_s32(0); + float32x4_t c00 = vdupq_n_f32(0.f); + float32x4_t c01 = vdupq_n_f32(0.f); + float32x4_t c02 = vdupq_n_f32(0.f); + float32x4_t c10 = vdupq_n_f32(0.f); + float32x4_t c11 = vdupq_n_f32(0.f); + float32x4_t c12 = vdupq_n_f32(0.f); + float32x4_t c20 = vdupq_n_f32(0.f); + float32x4_t c21 = vdupq_n_f32(0.f); + float32x4_t c22 = vdupq_n_f32(0.f); + const TA *Ap0 = A + lda * (i + 0); + const TA *Ap1 = A + lda * (i + 1); + const TA *Ap2 = A + lda * (i + 2); + const block_q8_0 *Bp0 = B + ldb * (j + 0); + const block_q8_0 *Bp1 = B + ldb * (j + 1); + const block_q8_0 *Bp2 = B + ldb * (j + 2); + for (int l = 0; l < k; ++l) { + c00 = vmlaq_n_f32( + c00, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap0 + l), load_lo(Bp0 + l)), + load_hi(Ap0 + l), load_hi(Bp0 + l))), + unhalf(Ap0[l].d) * unhalf(Bp0[l].d)); + c01 = vmlaq_n_f32( + c01, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap0 + l), load_lo(Bp1 + l)), + load_hi(Ap0 + l), load_hi(Bp1 + l))), + unhalf(Ap0[l].d) * unhalf(Bp1[l].d)); + c02 = vmlaq_n_f32( + c02, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap0 + l), load_lo(Bp2 + l)), + load_hi(Ap0 + l), load_hi(Bp2 + l))), + unhalf(Ap0[l].d) * unhalf(Bp2[l].d)); + c10 = vmlaq_n_f32( + c10, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap1 + l), load_lo(Bp0 + l)), + load_hi(Ap1 + l), load_hi(Bp0 + l))), + unhalf(Ap1[l].d) * unhalf(Bp0[l].d)); + c11 = vmlaq_n_f32( + c11, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap1 + l), load_lo(Bp1 + l)), + load_hi(Ap1 + l), load_hi(Bp1 + l))), + unhalf(Ap1[l].d) * unhalf(Bp1[l].d)); + c12 = vmlaq_n_f32( + c12, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap1 + l), load_lo(Bp2 + l)), + load_hi(Ap1 + l), load_hi(Bp2 + l))), + unhalf(Ap1[l].d) * unhalf(Bp2[l].d)); + c20 = vmlaq_n_f32( + c20, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap2 + l), load_lo(Bp0 + l)), + load_hi(Ap2 + l), load_hi(Bp0 + l))), + unhalf(Ap2[l].d) * unhalf(Bp0[l].d)); + c21 = vmlaq_n_f32( + c21, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap2 + l), load_lo(Bp1 + l)), + load_hi(Ap2 + l), load_hi(Bp1 + l))), + unhalf(Ap2[l].d) * unhalf(Bp1[l].d)); + c22 = vmlaq_n_f32( + c22, + vcvtq_f32_s32(vdotq_s32(vdotq_s32(zero, load_lo(Ap2 + l), load_lo(Bp2 + l)), + load_hi(Ap2 + l), load_hi(Bp2 + l))), + unhalf(Ap2[l].d) * unhalf(Bp2[l].d)); + } + C[ldc * (j + 0) + (i + 0)] = hsum(c00); + C[ldc * (j + 0) + (i + 1)] = hsum(c10); + C[ldc * (j + 0) + (i + 2)] = hsum(c20); + C[ldc * (j + 1) + (i + 0)] = hsum(c01); + C[ldc * (j + 1) + (i + 1)] = hsum(c11); + C[ldc * (j + 1) + (i + 2)] = hsum(c21); + C[ldc * (j + 2) + (i + 0)] = hsum(c02); + C[ldc * (j + 2) + (i + 1)] = hsum(c12); + C[ldc * (j + 2) + (i + 2)] = hsum(c22); + END_KERNEL() + } + + NOINLINE void gemm1x1(int m0, int m, int n0, int n) { + BEGIN_KERNEL(1, 1) + float32x4_t acc = vdupq_n_f32(0.f); + const TA *Ap = A + lda * i; + const block_q8_0 *Bp = B + ldb * j; + for (int l = 0; l < k; ++l) { + acc = vmlaq_n_f32(acc, + vcvtq_f32_s32(vdotq_s32( + vdotq_s32(vdupq_n_s32(0), load_lo(Ap + l), load_lo(Bp + l)), + load_hi(Ap + l), load_hi(Bp + l))), + unhalf(Ap[l].d) * unhalf(Bp[l].d)); + } + C[ldc * j + i] = hsum(acc); + END_KERNEL() + } + + inline int8x16_t load_lo(const block_q8_0 *b) { + return vld1q_s8(b->qs); + } + inline int8x16_t load_hi(const block_q8_0 *b) { + return vld1q_s8(b->qs + 16); + } + + inline int8x16_t load_lo(const block_q4_0 *b) { + return vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vld1q_u8(b->qs), + vdupq_n_u8(0x0f))), + vdupq_n_s8(0x8)); + } + inline int8x16_t load_hi(const block_q4_0 *b) { + return vsubq_s8(vreinterpretq_s8_u8(vshrq_n_u8(vld1q_u8(b->qs), 4)), + vdupq_n_s8(0x8)); + } + + const TA *const A; + const block_q8_0 *const B; + float *const C; + const int k; + const int lda; + const int ldb; + const int ldc; + const int ith; + const int nth; +}; +#endif // __ARM_FEATURE_DOTPROD + +#if defined(__AVX2__) || defined(__AVX512F__) +template <typename TA, typename TB, typename TC> +class tinyBLAS_Q0_AVX2 { + public: + tinyBLAS_Q0_AVX2(int k, + const TA *A, int lda, + const TB *B, int ldb, + TC *C, int ldc, + int ith, int nth) + : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) { + } + + void matmul(int m, int n, int task) { + if (task == GGML_TASK_TYPE_COMPUTE) + mnpack(0, m, 0, n); + } + + private: + NOINLINE void mnpack(int m0, int m, int n0, int n) { + int mc, nc, mp, np; + if (m - m0 <= 0 || n - n0 <= 0) + return; + if (m - m0 >= 4 && n - n0 >= 3) { + mc = 4; + nc = 3; + gemm4x3(m0, m, n0, n); + } else if (m - m0 >= 4 && n - n0 >= 1) { + mc = 4; + nc = 1; + gemm4x1(m0, m, n0, n); + } else if (m - m0 >= 1 && n - n0 >= 4) { + mc = 1; + nc = 4; + gemm1x4(m0, m, n0, n); + } else { + mc = 1; + nc = 1; + gemm1x1(m0, m, n0, n); + } + mp = m0 + (m - m0) / mc * mc; + np = n0 + (n - n0) / nc * nc; + mnpack(mp, m, n0, np); + mnpack(m0, mp, np, n); + mnpack(mp, m, np, n); + } + + NOINLINE void gemm4x3(int m0, int m, int n0, int n) { + BEGIN_KERNEL(4, 3) + __m256 c00 = _mm256_setzero_ps(); + __m256 c10 = _mm256_setzero_ps(); + __m256 c20 = _mm256_setzero_ps(); + __m256 c30 = _mm256_setzero_ps(); + __m256 c01 = _mm256_setzero_ps(); + __m256 c11 = _mm256_setzero_ps(); + __m256 c21 = _mm256_setzero_ps(); + __m256 c31 = _mm256_setzero_ps(); + __m256 c02 = _mm256_setzero_ps(); + __m256 c12 = _mm256_setzero_ps(); + __m256 c22 = _mm256_setzero_ps(); + __m256 c32 = _mm256_setzero_ps(); + const TA *Ap0 = A + lda * (i + 0); + const TA *Ap1 = A + lda * (i + 1); + const TA *Ap2 = A + lda * (i + 2); + const TA *Ap3 = A + lda * (i + 3); + const TB *Bp0 = B + ldb * (j + 0); + const TB *Bp1 = B + ldb * (j + 1); + const TB *Bp2 = B + ldb * (j + 2); + for (int l = 0; l < k; ++l) { + float da0 = unhalf(Ap0[l].d); + float da1 = unhalf(Ap1[l].d); + float da2 = unhalf(Ap2[l].d); + float da3 = unhalf(Ap3[l].d); + __m256i e0 = load(Ap0 + l); + __m256i e1 = load(Ap1 + l); + __m256i e2 = load(Ap2 + l); + __m256i e3 = load(Ap3 + l); + float db0 = unhalf(Bp0[l].d); + __m256 d00 = _mm256_set1_ps(da0 * db0); + __m256 d10 = _mm256_set1_ps(da1 * db0); + __m256 d20 = _mm256_set1_ps(da2 * db0); + __m256 d30 = _mm256_set1_ps(da3 * db0); + __m256i f0 = load(Bp0 + l); + __m256i u0 = _mm256_sign_epi8(f0, f0); + __m256i s00 = _mm256_sign_epi8(e0, f0); + __m256i s10 = _mm256_sign_epi8(e1, f0); + __m256i s20 = _mm256_sign_epi8(e2, f0); + __m256i s30 = _mm256_sign_epi8(e3, f0); + c00 = madd(d00, updot(u0, s00), c00); + c10 = madd(d10, updot(u0, s10), c10); + c20 = madd(d20, updot(u0, s20), c20); + c30 = madd(d30, updot(u0, s30), c30); + float db1 = unhalf(Bp1[l].d); + __m256 d01 = _mm256_set1_ps(da0 * db1); + __m256 d11 = _mm256_set1_ps(da1 * db1); + __m256 d21 = _mm256_set1_ps(da2 * db1); + __m256 d31 = _mm256_set1_ps(da3 * db1); + __m256i f1 = load(Bp1 + l); + __m256i u1 = _mm256_sign_epi8(f1, f1); + __m256i s01 = _mm256_sign_epi8(e0, f1); + __m256i s11 = _mm256_sign_epi8(e1, f1); + __m256i s21 = _mm256_sign_epi8(e2, f1); + __m256i s31 = _mm256_sign_epi8(e3, f1); + c01 = madd(d01, updot(u1, s01), c01); + c11 = madd(d11, updot(u1, s11), c11); + c21 = madd(d21, updot(u1, s21), c21); + c31 = madd(d31, updot(u1, s31), c31); + float db2 = unhalf(Bp2[l].d); + __m256 d02 = _mm256_set1_ps(da0 * db2); + __m256 d12 = _mm256_set1_ps(da1 * db2); + __m256 d22 = _mm256_set1_ps(da2 * db2); + __m256 d32 = _mm256_set1_ps(da3 * db2); + __m256i f2 = load(Bp2 + l); + __m256i u2 = _mm256_sign_epi8(f2, f2); + __m256i s02 = _mm256_sign_epi8(e0, f2); + __m256i s12 = _mm256_sign_epi8(e1, f2); + __m256i s22 = _mm256_sign_epi8(e2, f2); + __m256i s32 = _mm256_sign_epi8(e3, f2); + c02 = madd(d02, updot(u2, s02), c02); + c12 = madd(d12, updot(u2, s12), c12); + c22 = madd(d22, updot(u2, s22), c22); + c32 = madd(d32, updot(u2, s32), c32); + } + C[ldc * (j + 0) + (i + 0)] = hsum(c00); + C[ldc * (j + 0) + (i + 1)] = hsum(c10); + C[ldc * (j + 0) + (i + 2)] = hsum(c20); + C[ldc * (j + 0) + (i + 3)] = hsum(c30); + C[ldc * (j + 1) + (i + 0)] = hsum(c01); + C[ldc * (j + 1) + (i + 1)] = hsum(c11); + C[ldc * (j + 1) + (i + 2)] = hsum(c21); + C[ldc * (j + 1) + (i + 3)] = hsum(c31); + C[ldc * (j + 2) + (i + 0)] = hsum(c02); + C[ldc * (j + 2) + (i + 1)] = hsum(c12); + C[ldc * (j + 2) + (i + 2)] = hsum(c22); + C[ldc * (j + 2) + (i + 3)] = hsum(c32); + END_KERNEL() + } + + NOINLINE void gemm4x1(int m0, int m, int n0, int n) { + BEGIN_KERNEL(4, 1) + __m256 c0 = _mm256_setzero_ps(); + __m256 c1 = _mm256_setzero_ps(); + __m256 c2 = _mm256_setzero_ps(); + __m256 c3 = _mm256_setzero_ps(); + const TA *Ap0 = A + lda * (i + 0); + const TA *Ap1 = A + lda * (i + 1); + const TA *Ap2 = A + lda * (i + 2); + const TA *Ap3 = A + lda * (i + 3); + const TB *Bp = B + ldb * j; + for (int l = 0; l < k; ++l) { + float db0 = unhalf(Bp[l].d); + __m256i f = load(Bp + l); + __m256i u = _mm256_sign_epi8(f, f); + __m256 d0 = _mm256_set1_ps(unhalf(Ap0[l].d) * db0); + __m256 d1 = _mm256_set1_ps(unhalf(Ap1[l].d) * db0); + __m256 d2 = _mm256_set1_ps(unhalf(Ap2[l].d) * db0); + __m256 d3 = _mm256_set1_ps(unhalf(Ap3[l].d) * db0); + __m256i e0 = load(Ap0 + l); + __m256i e1 = load(Ap1 + l); + __m256i e2 = load(Ap2 + l); + __m256i e3 = load(Ap3 + l); + __m256i s0 = _mm256_sign_epi8(e0, f); + __m256i s1 = _mm256_sign_epi8(e1, f); + __m256i s2 = _mm256_sign_epi8(e2, f); + __m256i s3 = _mm256_sign_epi8(e3, f); + __m256 g0 = updot(u, s0); + __m256 g1 = updot(u, s1); + __m256 g2 = updot(u, s2); + __m256 g3 = updot(u, s3); + c0 = madd(d0, g0, c0); + c1 = madd(d1, g1, c1); + c2 = madd(d2, g2, c2); + c3 = madd(d3, g3, c3); + } + C[ldc * j + (i + 0)] = hsum(c0); + C[ldc * j + (i + 1)] = hsum(c1); + C[ldc * j + (i + 2)] = hsum(c2); + C[ldc * j + (i + 3)] = hsum(c3); + END_KERNEL() + } + + NOINLINE void gemm1x4(int m0, int m, int n0, int n) { + BEGIN_KERNEL(1, 4) + __m256 c0 = _mm256_setzero_ps(); + __m256 c1 = _mm256_setzero_ps(); + __m256 c2 = _mm256_setzero_ps(); + __m256 c3 = _mm256_setzero_ps(); + const TB *Bp0 = B + ldb * (j + 0); + const TB *Bp1 = B + ldb * (j + 1); + const TB *Bp2 = B + ldb * (j + 2); + const TB *Bp3 = B + ldb * (j + 3); + const TA *Ap = A + lda * i; + for (int l = 0; l < k; ++l) { + float da0 = unhalf(Ap[l].d); + __m256i f = load(Ap + l); + __m256i u = _mm256_sign_epi8(f, f); + __m256 d0 = _mm256_set1_ps(unhalf(Bp0[l].d) * da0); + __m256 d1 = _mm256_set1_ps(unhalf(Bp1[l].d) * da0); + __m256 d2 = _mm256_set1_ps(unhalf(Bp2[l].d) * da0); + __m256 d3 = _mm256_set1_ps(unhalf(Bp3[l].d) * da0); + __m256 g0 = updot(u, _mm256_sign_epi8(load(Bp0 + l), f)); + __m256 g1 = updot(u, _mm256_sign_epi8(load(Bp1 + l), f)); + __m256 g2 = updot(u, _mm256_sign_epi8(load(Bp2 + l), f)); + __m256 g3 = updot(u, _mm256_sign_epi8(load(Bp3 + l), f)); + c0 = madd(d0, g0, c0); + c1 = madd(d1, g1, c1); + c2 = madd(d2, g2, c2); + c3 = madd(d3, g3, c3); + } + C[ldc * (j + 0) + i] = hsum(c0); + C[ldc * (j + 1) + i] = hsum(c1); + C[ldc * (j + 2) + i] = hsum(c2); + C[ldc * (j + 3) + i] = hsum(c3); + END_KERNEL() + } + + NOINLINE void gemm1x1(int m0, int m, int n0, int n) { + BEGIN_KERNEL(1, 1) + __m256 c = _mm256_setzero_ps(); + const TA *Ap = A + lda * i; + const TB *Bp = B + ldb * j; + for (int l = 0; l < k; ++l) { + __m256 d = _mm256_set1_ps(unhalf(Ap[l].d) * unhalf(Bp[l].d)); + __m256i e = load(Ap + l); + __m256i f = load(Bp + l); + __m256 g = updot(_mm256_sign_epi8(e, e), _mm256_sign_epi8(f, e)); + c = madd(d, g, c); + } + C[ldc * j + i] = hsum(c); + END_KERNEL() + } + + inline __m256i load(const block_q8_0 *b) { + return _mm256_loadu_si256((const __m256i *)b->qs); + } + + inline __m256i load(const block_q4_0 *b) { + return _mm256_sub_epi8(denibble(b->qs), _mm256_set1_epi8(8)); + } + + inline __m256 updot(__m256i u, __m256i s) { + __m256i res; +#if defined(__AVXVNNI__) || (defined(__AVX512VNNI__) && defined(__AVX512VL__)) + res = _mm256_dpbusd_epi32(_mm256_setzero_si256(), u, s); +#else + res = _mm256_madd_epi16(_mm256_set1_epi16(1), _mm256_maddubs_epi16(u, s)); +#endif + return _mm256_cvtepi32_ps(res); + } + + static inline __m256i denibble(const uint8_t *p) { + const __m128i tmp = _mm_loadu_si128((const __m128i *)p); + const __m256i bytes = MM256_SET_M128I(_mm_srli_epi16(tmp, 4), tmp); + const __m256i lowMask = _mm256_set1_epi8(15); + return _mm256_and_si256(lowMask, bytes); + } + + const TA *const A; + const TB *const B; + TC *const C; + const int k; + const int lda; + const int ldb; + const int ldc; + const int ith; + const int nth; +}; +#endif // __AVX2__ + +} // namespace + +/** + * Performs optimized matrix multiplication on CPU. + * + * This subroutine may compute C = Aᵀ * B with column major ordering. + * Despite its name, this isn't a generalized implementation. Work is + * only performed when a handwritten kernel is written and available. + * Otherwise the caller should fall back to a general matmul routine. + * + * For example, for single-threaded single-precision GEMM you can say + * + * llamafile_sgemm(m, n, k, A, lda, B, ldb, C, ldc, + * 0, 1, GGML_TASK_TYPE_COMPUTE, + * GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32); + * + * @param m is rows in `A` and `C` + * @param n is cols in `B` and `C` + * @param k is cols in `A` and rows in `B` + * @param A is first input matrix (always transposed) + * @param lda is row stride of `A` + * @param B is second input matrix (never transposed) + * @param ldb is row stride of `B` + * @param C is input/output array of output matrices + * @param ldc is row stride of `C` + * @param ith is thread id (must be less than `nth`) + * @param nth is number of threads (must be greater than zero) + * @param task is GGML task type + * @param Atype is GGML data type of `A` + * @param Btype is GGML data type of `B` + * @param Ctype is GGML data type of `C` + * @return true if this function was able to service the matmul request + */ +bool llamafile_sgemm(int m, int n, int k, const void *A, int lda, const void *B, int ldb, void *C, + int ldc, int ith, int nth, int task, int Atype, int Btype, int Ctype) { + + assert(m >= 0); + assert(n >= 0); + assert(k >= 0); + assert(lda >= k); + assert(ldb >= k); + assert(ldc >= m); + assert(nth > 0); + assert(ith < nth); + assert(1ll * lda * m <= 0x7fffffff); + assert(1ll * ldb * n <= 0x7fffffff); + assert(1ll * ldc * n <= 0x7fffffff); + + if (Ctype != GGML_TYPE_F32) + return false; + + switch (Atype) { + + case GGML_TYPE_F32: { + if (Btype != GGML_TYPE_F32) + return false; +#if defined(__AVX512F__) + if (k % 16) + return false; + tinyBLAS<16, __m512, __m512, float, float, float> tb{ + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#elif defined(__AVX__) || defined(__AVX2__) + if (k % 8) + return false; + tinyBLAS<8, __m256, __m256, float, float, float> tb{ + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#elif defined(__ARM_NEON) + if (n < 4) + return false; + if (k % 4) + return false; + tinyBLAS<4, float32x4_t, float32x4_t, float, float, float> tb{ + k, (const float *)A, lda, + (const float *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#else + return false; +#endif + } + + case GGML_TYPE_F16: { +#if defined(__AVX512F__) + if (k % 16) + return false; + if (Btype != GGML_TYPE_F32) + return false; + tinyBLAS<16, __m512, __m512, ggml_fp16_t, float, float> tb{ + k, (const ggml_fp16_t *)A, lda, + (const float *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#elif (defined(__AVX__) || defined(__AVX2__)) && defined(__F16C__) + if (k % 8) + return false; + if (Btype != GGML_TYPE_F32) + return false; + tinyBLAS<8, __m256, __m256, ggml_fp16_t, float, float> tb{ + k, (const ggml_fp16_t *)A, lda, + (const float *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#elif defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) && !defined(_MSC_VER) + if (n < 8) + return false; + if (k % 8) + return false; + if (Btype != GGML_TYPE_F16) + return false; + tinyBLAS<8, float16x8_t, float16x8_t, ggml_fp16_t, ggml_fp16_t, float> tb{ + k, (const ggml_fp16_t *)A, lda, + (const ggml_fp16_t *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#elif defined(__ARM_NEON) && !defined(_MSC_VER) + if (k % 4) + return false; + if (Btype != GGML_TYPE_F32) + return false; + tinyBLAS<4, float32x4_t, float32x4_t, ggml_fp16_t, float, float> tb{ + k, (const ggml_fp16_t *)A, lda, + (const float *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#else + return false; +#endif + } + + case GGML_TYPE_Q8_0: { + if (Btype != GGML_TYPE_Q8_0) + return false; +#if defined(__AVX2__) || defined(__AVX512F__) + tinyBLAS_Q0_AVX2<block_q8_0, block_q8_0, float> tb{ + k, (const block_q8_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#elif defined(__ARM_FEATURE_DOTPROD) + tinyBLAS_Q0_ARM<block_q8_0> tb{ + k, (const block_q8_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#else + return false; +#endif + } + + case GGML_TYPE_Q4_0: { + if (Btype != GGML_TYPE_Q8_0) + return false; +#if defined(__AVX2__) || defined(__AVX512F__) + tinyBLAS_Q0_AVX2<block_q4_0, block_q8_0, float> tb{ + k, (const block_q4_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#elif defined(__ARM_FEATURE_DOTPROD) + tinyBLAS_Q0_ARM<block_q4_0> tb{ + k, (const block_q4_0 *)A, lda, + (const block_q8_0 *)B, ldb, + (float *)C, ldc, + ith, nth}; + tb.matmul(m, n, task); + return true; +#else + return false; +#endif + } + + default: + return false; + } + + (void)m; + (void)n; + (void)k; + (void)A; + (void)lda; + (void)B; + (void)ldb; + (void)C; + (void)ldc; + (void)ith; + (void)nth; + (void)task; + (void)Atype; + (void)Btype; + (void)Ctype; +} |