From 8cc91dc63c0df397d644a581b2cbeea74eb51ae0 Mon Sep 17 00:00:00 2001 From: Justine Tunney Date: Tue, 16 Apr 2024 14:55:30 -0400 Subject: [PATCH] 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/ --- CMakeLists.txt | 2 + Makefile | 10 +- Package.swift | 1 + build.zig | 15 +- common/common.cpp | 73 ++ common/common.h | 3 +- examples/llama-bench/llama-bench.cpp | 2 +- ggml-impl.h | 2 +- ggml-quants.c | 2 +- ggml.c | 54 ++ sgemm.cpp | 1148 ++++++++++++++++++++++++++ sgemm.h | 12 + 12 files changed, 1312 insertions(+), 12 deletions(-) create mode 100644 sgemm.cpp create mode 100644 sgemm.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 19fdfa46c..158174c20 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1151,6 +1151,8 @@ add_library(ggml OBJECT ggml-backend.h ggml-quants.c ggml-quants.h + sgemm.cpp + sgemm.h ${GGML_SOURCES_CUDA} ${GGML_HEADERS_CUDA} ${GGML_SOURCES_OPENCL} ${GGML_HEADERS_OPENCL} ${GGML_SOURCES_METAL} ${GGML_HEADERS_METAL} diff --git a/Makefile b/Makefile index 8f3e17da4..928fb14ca 100644 --- a/Makefile +++ b/Makefile @@ -219,6 +219,11 @@ ifdef LLAMA_DISABLE_LOGS MK_CPPFLAGS += -DLOG_DISABLE_LOGS endif # LLAMA_DISABLE_LOGS +# disable ggml.c's use of sgemm.cpp +ifdef LLAMA_NO_LLAMAFILE + MK_CPPFLAGS += -DGGML_USE_LLAMAFILE=0 +endif + # warnings WARN_FLAGS = -Wall -Wextra -Wpedantic -Wcast-qual -Wno-unused-function MK_CFLAGS += $(WARN_FLAGS) -Wshadow -Wstrict-prototypes -Wpointer-arith -Wmissing-prototypes -Werror=implicit-int \ @@ -676,13 +681,16 @@ ggml-backend.o: ggml-backend.c ggml.h ggml-backend.h ggml-quants.o: ggml-quants.c ggml.h ggml-quants.h ggml-common.h $(CC) $(CFLAGS) -c $< -o $@ +sgemm.o: sgemm.cpp sgemm.h ggml.h + $(CXX) $(CXXFLAGS) -c $< -o $@ + unicode.o: unicode.cpp unicode.h $(CXX) $(CXXFLAGS) -c $< -o $@ unicode-data.o: unicode-data.cpp unicode-data.h $(CXX) $(CXXFLAGS) -c $< -o $@ -OBJS += ggml-alloc.o ggml-backend.o ggml-quants.o unicode.o unicode-data.o +OBJS += ggml-alloc.o ggml-backend.o ggml-quants.o unicode.o unicode-data.o sgemm.o llama.o: llama.cpp unicode.h ggml.h ggml-alloc.h ggml-backend.h ggml-cuda.h ggml-metal.h llama.h $(CXX) $(CXXFLAGS) -c $< -o $@ diff --git a/Package.swift b/Package.swift index fbd0973be..183e64757 100644 --- a/Package.swift +++ b/Package.swift @@ -4,6 +4,7 @@ import PackageDescription var sources = [ "ggml.c", + "sgemm.cpp", "llama.cpp", "unicode.cpp", "unicode-data.cpp", diff --git a/build.zig b/build.zig index e05ca2120..c35e801f8 100644 --- a/build.zig +++ b/build.zig @@ -112,6 +112,7 @@ pub fn build(b: *std.build.Builder) !void { make.enable_lto = b.option(bool, "lto", "Enable LTO optimization, (default: false)") orelse false; const ggml = make.obj("ggml", "ggml.c"); + const sgemm = make.obj("sgemm", "sgemm.cpp"); const ggml_alloc = make.obj("ggml-alloc", "ggml-alloc.c"); const ggml_backend = make.obj("ggml-backend", "ggml-backend.c"); const ggml_quants = make.obj("ggml-quants", "ggml-quants.c"); @@ -128,14 +129,14 @@ pub fn build(b: *std.build.Builder) !void { const clip = make.obj("clip", "examples/llava/clip.cpp"); const llava = make.obj("llava", "examples/llava/llava.cpp"); - _ = make.exe("main", "examples/main/main.cpp", &.{ ggml, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo, sampling, console, grammar_parser }); - _ = make.exe("quantize", "examples/quantize/quantize.cpp", &.{ ggml, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo }); - _ = make.exe("perplexity", "examples/perplexity/perplexity.cpp", &.{ ggml, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo }); - _ = make.exe("embedding", "examples/embedding/embedding.cpp", &.{ ggml, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo }); - _ = make.exe("finetune", "examples/finetune/finetune.cpp", &.{ ggml, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo, train }); - _ = make.exe("train-text-from-scratch", "examples/train-text-from-scratch/train-text-from-scratch.cpp", &.{ ggml, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo, train }); + _ = make.exe("main", "examples/main/main.cpp", &.{ ggml, sgemm, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo, sampling, console, grammar_parser }); + _ = make.exe("quantize", "examples/quantize/quantize.cpp", &.{ ggml, sgemm, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo }); + _ = make.exe("perplexity", "examples/perplexity/perplexity.cpp", &.{ ggml, sgemm, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo }); + _ = make.exe("embedding", "examples/embedding/embedding.cpp", &.{ ggml, sgemm, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo }); + _ = make.exe("finetune", "examples/finetune/finetune.cpp", &.{ ggml, sgemm, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo, train }); + _ = make.exe("train-text-from-scratch", "examples/train-text-from-scratch/train-text-from-scratch.cpp", &.{ ggml, sgemm, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo, train }); - const server = make.exe("server", "examples/server/server.cpp", &.{ ggml, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo, sampling, grammar_parser, clip, llava }); + const server = make.exe("server", "examples/server/server.cpp", &.{ ggml, sgemm, ggml_alloc, ggml_backend, ggml_quants, llama, unicode, unicode_data, common, json_schema_to_grammar, buildinfo, sampling, grammar_parser, clip, llava }); if (server.target.isWindows()) { server.linkSystemLibrary("ws2_32"); } diff --git a/common/common.cpp b/common/common.cpp index 52576cba3..cf69535e2 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -108,6 +108,79 @@ int32_t get_num_physical_cores() { return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4; } +#if defined(__x86_64__) && defined(__linux__) +#include + +static void cpuid(unsigned leaf, unsigned subleaf, + unsigned *eax, unsigned *ebx, unsigned *ecx, unsigned *edx) { + __asm__("movq\t%%rbx,%%rsi\n\t" + "cpuid\n\t" + "xchgq\t%%rbx,%%rsi" + : "=a"(*eax), "=S"(*ebx), "=c"(*ecx), "=d"(*edx) + : "0"(leaf), "2"(subleaf)); +} + +static int pin_cpu(int cpu) { + cpu_set_t mask; + CPU_ZERO(&mask); + CPU_SET(cpu, &mask); + return pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask); +} + +static bool is_hybrid_cpu(void) { + unsigned eax, ebx, ecx, edx; + cpuid(7, 0, &eax, &ebx, &ecx, &edx); + return !!(edx & (1u << 15)); +} + +static bool is_running_on_efficiency_core(void) { + unsigned eax, ebx, ecx, edx; + cpuid(0x1a, 0, &eax, &ebx, &ecx, &edx); + int intel_atom = 0x20; + int core_type = (eax & 0xff000000u) >> 24; + return core_type == intel_atom; +} + +static int count_math_cpus(int cpu_count) { + int result = 0; + for (int cpu = 0; cpu < cpu_count; ++cpu) { + if (pin_cpu(cpu)) { + return -1; + } + if (is_running_on_efficiency_core()) { + continue; // efficiency cores harm lockstep threading + } + ++cpu; // hyperthreading isn't useful for linear algebra + ++result; + } + return result; +} + +#endif // __x86_64__ && __linux__ + +/** + * Returns number of CPUs on system that are useful for math. + */ +int get_math_cpu_count() { +#if defined(__x86_64__) && defined(__linux__) + int cpu_count = sysconf(_SC_NPROCESSORS_ONLN); + if (cpu_count < 1) { + return get_num_physical_cores(); + } + if (is_hybrid_cpu()) { + cpu_set_t affinity; + if (!pthread_getaffinity_np(pthread_self(), sizeof(affinity), &affinity)) { + int result = count_math_cpus(cpu_count); + pthread_setaffinity_np(pthread_self(), sizeof(affinity), &affinity); + if (result > 0) { + return result; + } + } + } +#endif + return get_num_physical_cores(); +} + void process_escapes(std::string & input) { std::size_t input_len = input.length(); std::size_t output_idx = 0; diff --git a/common/common.h b/common/common.h index 65272b0ba..cca44268e 100644 --- a/common/common.h +++ b/common/common.h @@ -39,6 +39,7 @@ extern char const *LLAMA_BUILD_TARGET; struct llama_control_vector_load_info; +int get_math_cpu_count(); int32_t get_num_physical_cores(); // @@ -48,7 +49,7 @@ int32_t get_num_physical_cores(); struct gpt_params { uint32_t seed = LLAMA_DEFAULT_SEED; // RNG seed - int32_t n_threads = get_num_physical_cores(); + int32_t n_threads = get_math_cpu_count(); int32_t n_threads_draft = -1; int32_t n_threads_batch = -1; // number of threads to use for batch processing (-1 = use n_threads) int32_t n_threads_batch_draft = -1; diff --git a/examples/llama-bench/llama-bench.cpp b/examples/llama-bench/llama-bench.cpp index 27e113203..8b532c8b6 100644 --- a/examples/llama-bench/llama-bench.cpp +++ b/examples/llama-bench/llama-bench.cpp @@ -190,7 +190,7 @@ static const cmd_params cmd_params_defaults = { /* n_ubatch */ {512}, /* type_k */ {GGML_TYPE_F16}, /* type_v */ {GGML_TYPE_F16}, - /* n_threads */ {get_num_physical_cores()}, + /* n_threads */ {get_math_cpu_count()}, /* n_gpu_layers */ {99}, /* split_mode */ {LLAMA_SPLIT_MODE_LAYER}, /* main_gpu */ {0}, diff --git a/ggml-impl.h b/ggml-impl.h index e68b72877..0c997d3ed 100644 --- a/ggml-impl.h +++ b/ggml-impl.h @@ -88,7 +88,7 @@ typedef uint16_t ggml_fp16_internal_t; #if defined(_MSC_VER) || defined(__MINGW32__) #include #else -#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__SSSE3__) || defined(__SSE3__) +#if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__) || defined(__SSSE3__) || defined(__SSE3__) || defined(__SSE__) #if !defined(__riscv) #include #endif diff --git a/ggml-quants.c b/ggml-quants.c index 32e84434a..32360a1f1 100644 --- a/ggml-quants.c +++ b/ggml-quants.c @@ -132,7 +132,7 @@ static inline __m256 sum_i16_pairs_float(const __m256i x) { } static inline __m256 mul_sum_us8_pairs_float(const __m256i ax, const __m256i sy) { -#if defined(__AVXVNNI__) || defined(__AVX512VNNI__) +#if defined(__AVXVNNI__) || (defined(__AVX512VNNI__) && defined(__AVX512VL__)) const __m256i zero = _mm256_setzero_si256(); const __m256i summed_pairs = _mm256_dpbusd_epi32(zero, ax, sy); return _mm256_cvtepi32_ps(summed_pairs); diff --git a/ggml.c b/ggml.c index 14288d29d..119686be6 100644 --- a/ggml.c +++ b/ggml.c @@ -4,6 +4,7 @@ #include "ggml-impl.h" #include "ggml-quants.h" #include "ggml.h" +#include "sgemm.h" #if defined(_MSC_VER) || defined(__MINGW32__) #include // using malloc.h with MSC/MINGW @@ -32,6 +33,14 @@ #include #endif +#ifndef GGML_USE_LLAMAFILE +#ifdef __ARM_FEATURE_MATMUL_INT8 +#define GGML_USE_LLAMAFILE 0 +#else +#define GGML_USE_LLAMAFILE 1 +#endif +#endif + #if defined(_MSC_VER) // disable "possible loss of data" to avoid hundreds of casts // we should just be careful :) @@ -10810,6 +10819,28 @@ static void ggml_compute_forward_mul_mat( } #endif +#if GGML_USE_LLAMAFILE + if (nb10 == ggml_type_size(src1->type)) { + for (int64_t i13 = 0; i13 < ne13; i13++) + for (int64_t i12 = 0; i12 < ne12; i12++) + if (!llamafile_sgemm(ne01, ne11, ne00/ggml_blck_size(src0->type), + (const char *)src0->data + i12/r2*nb02 + i13/r3*nb03, + nb01/ggml_type_size(src0->type), + (const char *)src1->data + i12*nb12 + i13*nb13, + nb11/ggml_type_size(src1->type), + (char *)dst->data + i12*nb2 + i13*nb3, + nb1/ggml_type_size(dst->type), + ith, nth, + params->type, + src0->type, + src1->type, + dst->type)) + goto UseGgmlGemm1; + return; + } +UseGgmlGemm1:; +#endif + if (params->type == GGML_TASK_TYPE_INIT) { if (ith != 0) { return; @@ -10841,6 +10872,29 @@ static void ggml_compute_forward_mul_mat( const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; const size_t row_size = ggml_row_size(vec_dot_type, ne10); +#if GGML_USE_LLAMAFILE + if (nb10 == ggml_type_size(src1->type) || src1->type != vec_dot_type) { + for (int64_t i13 = 0; i13 < ne13; i13++) + for (int64_t i12 = 0; i12 < ne12; i12++) + if (!llamafile_sgemm(ne01, ne11, ne00/ggml_blck_size(src0->type), + (const char *)src0->data + i12/r2*nb02 + i13/r3*nb03, + nb01/ggml_type_size(src0->type), + (const char *)wdata + (nb12/ggml_type_size(src1->type)*ggml_type_size(vec_dot_type)*i12 + + nb13/ggml_type_size(src1->type)*ggml_type_size(vec_dot_type)*i13), + row_size/ggml_type_size(vec_dot_type), + (char *)dst->data + i12*nb2 + i13*nb3, + nb1/ggml_type_size(dst->type), + ith, nth, + params->type, + src0->type, + vec_dot_type, + dst->type)) + goto UseGgmlGemm2; + return; + } +UseGgmlGemm2:; +#endif + const int64_t nr0 = ne01; // src0 rows const int64_t nr1 = ne1*ne12*ne13; // src1 rows diff --git a/sgemm.cpp b/sgemm.cpp new file mode 100644 index 000000000..6900f04cf --- /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 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 +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 +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 +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(B + ldb * (j + 0) + l); + V k1 = load(B + ldb * (j + 1) + l); + V k2 = load(B + ldb * (j + 2) + l); + V k3 = load(B + ldb * (j + 3) + l); + V k4 = load(B + ldb * (j + 4) + l); + V a0 = load(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(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(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(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(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(B + ldb * (j + 0) + l); + V k1 = load(B + ldb * (j + 1) + l); + V k2 = load(B + ldb * (j + 2) + l); + V k3 = load(B + ldb * (j + 3) + l); + V a0 = load(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(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(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(A + lda * (i + 0) + l); + c00 = madder(a, load(B + ldb * (j + 0) + l), c00, &e00); + c01 = madder(a, load(B + ldb * (j + 1) + l), c01, &e01); + c02 = madder(a, load(B + ldb * (j + 2) + l), c02, &e02); + c03 = madder(a, load(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(B + ldb * (j + 0) + l); + c00 = madder(load(A + lda * (i + 0) + l), b, c00, &e00); + c10 = madder(load(A + lda * (i + 1) + l), b, c10, &e10); + c20 = madder(load(A + lda * (i + 2) + l), b, c20, &e20); + c30 = madder(load(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(A + lda * i + l), + load(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 +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 +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 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 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 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 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; +} diff --git a/sgemm.h b/sgemm.h new file mode 100644 index 000000000..da23b209c --- /dev/null +++ b/sgemm.h @@ -0,0 +1,12 @@ +#pragma once +#include +#ifdef __cplusplus +extern "C" { +#endif + +bool llamafile_sgemm(int, int, int, const void *, int, const void *, int, + void *, int, int, int, int, int, int, int); + +#ifdef __cplusplus +} +#endif