root/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggerganov/llama.cpp.git synced 2024-11-11 21:39:52 +00:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity
e931888d50
llama.cpp/sgemm.cpp

1149 lines
40 KiB
C++
Raw Normal View History

ggml : add llamafile sgemm (#6414) This change upstreams llamafile's cpu matrix multiplication kernels which improve image and prompt evaluation speed. For starters, Q4_0 and Q8_0 weights should go ~40% faster on CPU. The biggest benefits are with data types like f16 / f32, which process prompts 2x faster thus making them faster than quantized data types for prompt evals. This change also introduces bona fide AVX512 support since tinyBLAS is able to exploit the larger register file. For example, on my CPU llama.cpp llava-cli processes an image prompt at 305 tokens/second, using the Q4_K and Q4_0 types, which has always been faster than if we used f16 LLaVA weights, which at HEAD go 188 tokens/second. With this change, f16 LLaVA performance leap frogs to 464 tokens/second. On Intel Core i9-14900K this change improves F16 prompt perf by 5x. For example, using llama.cpp at HEAD with Mistral 7b f16 to process a 215 token prompt will go 13 tok/sec. This change has fixes making it go 52 tok/sec. It's mostly thanks to my vectorized outer product kernels but also because I added support for correctly counting the number of cores on Alderlake, so the default thread count discounts Intel's new efficiency cores. Only Linux right now can count cores. This work was sponsored by Mozilla who's given permission to change the license of this code from Apache 2.0 to MIT. To read more about what's improved, and how it works, see: https://justine.lol/matmul/
2024-04-16 18:55:30 +00:00
// -*- 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;
}
Reference in New Issue Copy Permalink