opencoti F5-opt W2 (#290) — wholesale ik_llama.cpp CPU Flash-Attention engine import Vendors ik_llama.cpp's CPU Flash-Attention engine behind a build flag (GGML_IQK_FLASH_ATTENTION) plus a single surgical dispatch hook, giving the M2 CPU-half (HeadInfer's CPU-resident attention heads) a faster token-generation FA kernel. Off by default; --iqk-flash-attn on engages it. Flag-OFF is byte-identical to the pre-0040 binary: the hook early-returns when the process-global toggle is false, and the engine TUs are only reachable through that call site. opencoti-hook: f5-opt-cpufa — see docs/features/advanced_kv.md UPSTREAM PROVENANCE (for future re-sync) Source repo : https://github.com/ikawrakow/ik_llama.cpp Commit SHA : 8960c5ba5ee9db30ba838304373aa4dbec9f7cbd (HEAD, FA files 2026-05-29) License : MIT Author : Iwan Kawrakow (Copyright (c) 2024-2025 Iwan Kawrakow) Related PRs : ik_llama #332 (CPU FA TG threading); #208 (Q8_KV cache type — rides in with the engine, enabled separately for #291). Files vendored verbatim under llama.cpp/ggml/src/iqk/ (+ fa/), each carrying a top banner naming its upstream path; do not edit in place, re-sync from source: iqk/iqk_flash_attn.cpp ← ik_llama ggml/src/iqk/iqk_flash_attn.cpp iqk/iqk_fa_dispatch.cpp ← ik_llama ggml/src/iqk/iqk_fa_dispatch.cpp iqk/fa/iqk_fa_templates.h ← ik_llama ggml/src/iqk/fa/iqk_fa_templates.h iqk/fa/iqk_fa_{64_64,96_96,128_128,192_128,192_192,256_256,320_256,512_512,576_512}.cpp ← ik_llama ggml/src/iqk/fa/* (9 head-size instantiations) iqk/iqk_gemm_{floats,kquants,legacy_quants}.{cpp,h}, iqk_gemm_ktquants.h ← ik_llama ggml/src/iqk/* (FA dot-product backends) iqk/iqk_quantize.{cpp,h}, iqk_common.h, iqk_common_extra.h, iqk_config.h, iqk/iqk_utils.h, iqk_mul_mat.h, iqk_flash_impl.h, iqk_ggml_type_ext.h, iqk/ggml-common.h ← ik_llama ggml/src/iqk/* + ggml/src support headers OPENCOTI ADAPTATIONS (delta from upstream — the re-sync touch points) - llama.cpp/BUILD.mk: adds the engine TUs to the GGML C++ SRCS list and an IQK_FA_OBJS per-object rule granting them -Xx86_64-mavx2/-mf16c/-mfma -DGGML_IQK_FLASH_ATTENTION -Wno-enum-constexpr-conversion (bug-243: clang rejects the >63 ggml_type casts in the FA switch labels otherwise). ops.cpp.o and ggml-cpu.c.o also receive -DGGML_IQK_FLASH_ATTENTION for the two hooks. - llama.cpp/ggml/include/ggml-iqk-flash-attn.h (NEW): the public toggle (set/get) + 3 engine entry decls, pulling ONLY ggml.h. Core TUs include this, never iqk_mul_mat.h/iqk_config.h, whose NEON shim collides with ggml-cpu-impl.h on the aarch64 cosmocc slice (bug-244). - llama.cpp/ggml/src/ggml-cpu/ops.cpp: dispatch hook in ggml_compute_forward_flash_attn_ext_f16, gated GGML_IQK_FLASH_ATTENTION && __x86_64__ (bug-245: engine symbols exist only on the AVX2 x86_64 slice). Calls iqk_flash_attn_noalibi(); on false, falls through to the mainline _f16 path unchanged. - llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c: work-buffer sizing hook (MAX(cur, iqk_fa_work_buffer_size(...))) in the FLASH_ATTN_EXT work-size case. - llama.cpp/ggml/src/ggml-iqk-flash-attn.cpp (NEW): the atomic toggle TU. Intentionally EXCLUDED from IQK_FA_OBJS — it holds no engine/SIMD code and must not carry the enable guard. - llama.cpp/common/{arg.cpp,common.h}: --iqk-flash-attn on|off CLI (LLAMA_EXAMPLE_SERVER-gated, LLAMA_ARG_IQK_FLASH_ATTN) + params.iqk_flash_attn. VERIFICATION (2026-05-30, Gemma-4 A4B v6-coder Q4_K_M + Qwen2.5 0.5B, solidPC RTX 3090) - Flag-OFF off-path identity: neo-pipeline.bench C4 (--neo-pipeline off byte-identical) PASS; advanced-kv-stack 8/9 (sole fail C5 = pre-existing bug-207 NEO FP-noise, tracked by #280 — unrelated to this patch, ran iqk OFF). - Flag-ON quality (RULER vt + niah_single_1 @ 4k/32k, M2 + full F5 stack + --iqk-flash-attn on): 100% all 8 cells; ON vs OFF wall-time flat. The engine is correctness-clean; on these prefill-bound RULER cells the CPU-split residency (not CPU-FA kernel speed) is the perf lever — addressed by M7 Rolling KV / GPU_RESIDENT default (docs/features/rolling_kv.md Decision 4). Captured via snapshot-diff (bug-121): state-pre-0040 = post-0036 vendored tree (BUILD.mk/common.h/arg.cpp from .opencoti/state-post-0036; ops.cpp from state-post-0035; ggml-cpu.c from pristine HEAD — no prior patch touches it); state-post-0040 = this W2 tree. Diffs follow. diff -ruN state-pre-0040/llama.cpp/BUILD.mk state-post-0040/llama.cpp/BUILD.mk diff --git a/llama.cpp/BUILD.mk b/llama.cpp/BUILD.mk --- a/llama.cpp/BUILD.mk +++ b/llama.cpp/BUILD.mk @@ -41,7 +41,23 @@ GGML_SRCS_CPP := \ llama.cpp/ggml/src/ggml-cpu/unary-ops.cpp \ llama.cpp/ggml/src/ggml-cpu/vec.cpp \ llama.cpp/ggml/src/ggml-cpu/amx/amx.cpp \ - llama.cpp/ggml/src/ggml-cpu/amx/mmq.cpp + llama.cpp/ggml/src/ggml-cpu/amx/mmq.cpp \ + llama.cpp/ggml/src/ggml-iqk-flash-attn.cpp \ + llama.cpp/ggml/src/iqk/iqk_flash_attn.cpp \ + llama.cpp/ggml/src/iqk/iqk_fa_dispatch.cpp \ + llama.cpp/ggml/src/iqk/iqk_gemm_floats.cpp \ + llama.cpp/ggml/src/iqk/iqk_gemm_kquants.cpp \ + llama.cpp/ggml/src/iqk/iqk_gemm_legacy_quants.cpp \ + llama.cpp/ggml/src/iqk/iqk_quantize.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_64_64.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_96_96.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_128_128.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_192_128.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_192_192.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_256_256.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_320_256.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_512_512.cpp \ + llama.cpp/ggml/src/iqk/fa/iqk_fa_576_512.cpp GGML_OBJS := \ $(GGML_SRCS_C:%.c=o/$(MODE)/%.c.o) \ @@ -518,6 +534,55 @@ o/$(MODE)/llama.cpp/ggml/src/ggml-quants.c.o \ o/$(MODE)/llama.cpp/ggml/src/ggml-cpu/quants.c.o: \ private CCFLAGS += -O3 -mgcc +# opencoti F5-opt W2 (#290): ik_llama CPU Flash-Attention engine. +# The iqk engine TUs gate every SIMD path behind IQK_IMPLEMENT (= __AVX2__ || +# __ARM_FEATURE_DOTPROD). cosmocc's default x86_64 baseline lacks AVX2, so +# without these flags the whole engine compiles to empty stubs. Give the engine +# TUs AVX2/F16C/FMA on the x86_64 slice + the master enable GGML_IQK_FLASH_ATTENTION. +# (Shipped multi-host ARM/AVX-512 dispatch is a later concern; solidPC is AVX2.) +# The flag TU ggml-iqk-flash-attn.cpp is intentionally EXCLUDED — it is a plain +# atomic toggle with no engine/SIMD code and must not carry the enable guard. +IQK_FA_OBJS := \ + o/$(MODE)/llama.cpp/ggml/src/iqk/iqk_flash_attn.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/iqk_fa_dispatch.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/iqk_gemm_floats.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/iqk_gemm_kquants.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/iqk_gemm_legacy_quants.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/iqk_quantize.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_64_64.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_96_96.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_128_128.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_192_128.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_192_192.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_256_256.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_320_256.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_512_512.cpp.o \ + o/$(MODE)/llama.cpp/ggml/src/iqk/fa/iqk_fa_576_512.cpp.o + +# opencoti F5-opt W2 (#290): the iqk engine carries ik_llama.cpp's EXTENDED +# ggml_type tags (Q8_KV=151, Q8_0_R8=208, ...) supplied via the isolation header +# iqk_ggml_type_ext.h as `#define ((ggml_type)N)`. llamafile's real enum closes at +# GGML_TYPE_COUNT=42, so clang derives the range [0,63] and -Wenum-constexpr-conversion +# (a DEFAULT-ERROR diagnostic) rejects casting 151 etc. to ggml_type in the FA +# template switch labels + constexpr `type` members. The enum is int-backed at +# runtime so 151 is representable and codegen is correct; these tags are internal +# iqk dispatch keys that never reach llamafile core (its type_traits[COUNT] table is +# never indexed by them). We suppress the pedantic constant-expression check on the +# iqk TUs ONLY — keeping the hard-isolation invariant (ggml.h + GGML_TYPE_COUNT +# untouched) rather than poisoning the shared enum. +$(IQK_FA_OBJS): private CCFLAGS += \ + -Xx86_64-mavx2 -Xx86_64-mf16c -Xx86_64-mfma -DGGML_IQK_FLASH_ATTENTION \ + -Wno-enum-constexpr-conversion + +# opencoti F5-opt W2 (#290): the two core ggml-cpu TUs that carry the surgical +# dispatch hooks (ops.cpp = the FA dispatch; ggml-cpu.c = the wdata work-size +# bump) need GGML_IQK_FLASH_ATTENTION defined to compile the hook in. They stay +# at the -mgcc baseline (NOT -mavx2) — the hook code is a plain function call into +# the AVX2-compiled iqk engine, no SIMD of its own. Adding ONLY the -D keeps a +# no-flag build byte-identical (the hook #if-compiles out entirely). +o/$(MODE)/llama.cpp/ggml/src/ggml-cpu/ops.cpp.o: private CCFLAGS += -DGGML_IQK_FLASH_ATTENTION +o/$(MODE)/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c.o: private CCFLAGS += -DGGML_IQK_FLASH_ATTENTION + # ============================================================================== # Tool executables # ============================================================================== diff --git a/llama.cpp/common/arg.cpp b/llama.cpp/common/arg.cpp --- a/llama.cpp/common/arg.cpp +++ b/llama.cpp/common/arg.cpp @@ -9,6 +9,10 @@ #include "sampling.h" #include "speculative.h" #include "preset.h" +// opencoti F5-opt W2 (#290): the --iqk-flash-attn handler flips the ggml-cpu +// process-global that the FA dispatch hook reads. Header carries only the +// toggle getter/setter (no iqk_config.h drag — see bug-244). +#include "ggml-iqk-flash-attn.h" // fix problem with std::min and std::max #if defined(_WIN32) @@ -1440,6 +1444,19 @@ common_params_context common_params_parser_init(common_params & params, llama_ex else throw std::invalid_argument("--neo-pipeline must be off|on|auto"); } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_NEO_PIPELINE")); + // opencoti F5-opt W2 (#290) iqk CPU flash-attention — see docs/features/advanced_kv.md + add_opt(common_arg( + {"--iqk-flash-attn"}, "MODE", + "iqk CPU flash-attention (vendored ik_llama.cpp engine): off (default), on. " + "Speeds up the M2 CPU-half attention on the token-generation path; falls back " + "to the mainline kernel for unsupported (type, shape) cases. x86_64-only.", + [](common_params & params, const std::string & value) { + if (value == "off" || value == "false" || value == "0") params.iqk_flash_attn = false; + else if (value == "on" || value == "true" || value == "1") params.iqk_flash_attn = true; + else throw std::invalid_argument("--iqk-flash-attn must be on|off"); + ggml_iqk_flash_attn_set_enabled(params.iqk_flash_attn); + } + ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_IQK_FLASH_ATTN")); // opencoti F5-opt W1 (#293) pcie profile — see docs/features/rolling_kv.md add_opt(common_arg( {"--pcie-autodetect"}, "MODE", diff --git a/llama.cpp/common/common.h b/llama.cpp/common/common.h --- a/llama.cpp/common/common.h +++ b/llama.cpp/common/common.h @@ -450,6 +450,11 @@ struct common_params { // Rolling KV (#296) tile sizing. autodetect reads the rebar-probe JSON / // nvidia-smi; bw_gbps > 0 forces a manual override (0 = autodetect/default). bool pcie_autodetect = true; + // opencoti F5-opt W2 (#290) iqk CPU flash-attention — see docs/features/advanced_kv.md + // Engage the vendored ik_llama.cpp CPU Flash-Attention engine on the M2 + // CPU-half (token-generation path). Off by default; flag-off is byte-identical. + // x86_64-only at runtime (the engine compiles only on the AVX2 x86_64 slice). + bool iqk_flash_attn = false; float pcie_bw_gbps = 0.0f; int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) int32_t n_ubatch = 512; // physical batch size for prompt processing (must be >=32 to use BLAS) diff --git a/llama.cpp/ggml/include/ggml-iqk-flash-attn.h b/llama.cpp/ggml/include/ggml-iqk-flash-attn.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/include/ggml-iqk-flash-attn.h @@ -0,0 +1,57 @@ +// opencoti F5-opt W2 (#290) — ik_llama CPU Flash-Attention enable switch. +// +// A process-global on/off toggle for the vendored ik_llama.cpp CPU +// Flash-Attention engine (ggml/src/iqk/*). Mirrors the ggml-neo-pipeline +// global/setter pattern so the ggml-cpu dispatch hook in +// ggml_compute_forward_flash_attn_ext_f16 can read it WITHOUT a header +// dependency on llama/common. Default OFF: when off, the hook fast-skips and +// the FA op path is byte-identical to upstream. +// +// The flag is flipped once at server startup from the runtime tactic that +// resolves --iqk-flash-attn (common/arg.cpp). It is a static process-wide +// toggle (no per-layer state), so a single setter call at init suffices. +#pragma once + +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +// Enable / disable the iqk CPU Flash-Attention engine. Off by default. +GGML_API void ggml_iqk_flash_attn_set_enabled(bool enabled); +GGML_API bool ggml_iqk_flash_attn_enabled(void); + +// opencoti F5-opt W2 (#290): public entry points of the vendored ik_llama engine, +// re-declared here so the ggml-cpu dispatch hooks (ops.cpp + ggml-cpu.c) can call +// them WITHOUT including ggml/src/iqk/iqk_mul_mat.h — that header drags in +// iqk_config.h whose ARM ggml_vdotq_s32 NEON shim collides with the definition +// ggml-cpu-impl.h already provides in those TUs (bug-244). These signatures MUST +// stay byte-for-byte identical to ggml/src/iqk/iqk_mul_mat.h (extern "C", so the +// linker matches by name only — a silent mismatch would corrupt args). Re-sync +// both on any upstream ik_llama bump. +typedef void (*barrier_t)(void *); + +// Scratch (params->wdata) bytes the engine needs for one FA op at `nthread` threads. +size_t iqk_fa_work_buffer_size(const struct ggml_tensor * dst, int nthread); + +// Compute flash-attention for the non-alibi case. Returns true if it handled the +// op (result written to `qkv`); false if the (type, shape) case is unsupported or +// max_bias>0 (alibi) — caller must then run the mainline FA path. +bool iqk_flash_attn_noalibi(int type_q, int type_mask, float max_bias, + int neq3, int neq2, long nbq3, long nbq2, + int nek3, int nek2, long nbk3, long nbk2, + int nev3, int nev2, long nbv3, long nbv2, + int ne2, int ne1, long nb1, + int type_k, int type_v, int Dk, int Dv, + int nq, int nk, + int stride_q, int stride_k, int stride_v, int stride_m, + const void * q, const void * k, const void * v, + const void * mask, const void * sinks, + float scale, float softcap, float * qkv, + void * work_buffer, barrier_t barrier, void * barrier_data, + int ith, int nth, int n_swa); + +#ifdef __cplusplus +} +#endif diff --git a/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c b/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c --- a/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c +++ b/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c @@ -15,6 +15,15 @@ #include "ggml.h" #include "common.h" +#if defined(GGML_IQK_FLASH_ATTENTION) +// opencoti F5-opt W2 (#290): iqk_fa_work_buffer_size() — sizes the wdata scratch +// the vendored ik_llama CPU Flash-Attention engine needs (graph-plan side of the +// ops.cpp dispatch hook). Declared in ggml-iqk-flash-attn.h, NOT iqk/iqk_mul_mat.h, +// to avoid iqk_config.h's ARM ggml_vdotq_s32 shim colliding with ggml-cpu-impl.h +// (bug-244). Gated behind GGML_IQK_FLASH_ATTENTION (per-object on ggml-cpu.c.o). +#include "ggml-iqk-flash-attn.h" +#endif + #if defined(_MSC_VER) || defined(__MINGW32__) #include // using malloc.h with MSC/MINGW #elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) @@ -2969,6 +2978,16 @@ struct ggml_cplan ggml_graph_plan( size_t decode = sizeof(float)*(neq2*n_chunks*(2+DV) + n_tasks*(DK + 2*DV)); cur += MAX(prefill, decode); +#if defined(GGML_IQK_FLASH_ATTENTION) && defined(__x86_64__) + // opencoti F5-opt W2 (#290): reserve enough wdata for the + // vendored ik_llama CPU-FA engine so it never overruns when + // --iqk-flash-attn is on. Scratch sizing only — does not + // affect outputs, so flag-OFF decode stays byte-identical. + // Mirrors ik_llama's graph-plan reservation. x86_64-only: + // matches the dispatch hook in ops.cpp (the engine symbols + // only exist on the AVX2 x86_64 slice — see that hook). + cur = MAX(cur, iqk_fa_work_buffer_size(node, n_tasks)); +#endif } break; case GGML_OP_FLASH_ATTN_BACK: { diff --git a/llama.cpp/ggml/src/ggml-cpu/ops.cpp b/llama.cpp/ggml/src/ggml-cpu/ops.cpp --- a/llama.cpp/ggml/src/ggml-cpu/ops.cpp +++ b/llama.cpp/ggml/src/ggml-cpu/ops.cpp @@ -8,6 +8,17 @@ #include "unary-ops.h" #include "vec.h" +#if defined(GGML_IQK_FLASH_ATTENTION) +// opencoti F5-opt W2 (#290): vendored ik_llama.cpp CPU Flash-Attention engine. +// Gated entirely behind GGML_IQK_FLASH_ATTENTION (set per-object on ops.cpp.o in +// BUILD.mk) plus a runtime toggle, so a build without the flag is byte-identical. +// This header carries BOTH the runtime toggle getter AND the engine entry decls +// (iqk_flash_attn_noalibi/barrier_t) — we deliberately do NOT include +// iqk/iqk_mul_mat.h, which would drag in iqk_config.h's ARM ggml_vdotq_s32 shim +// and collide with ggml-cpu-impl.h (bug-244). +#include "ggml-iqk-flash-attn.h" +#endif + #include #include #include @@ -9068,6 +9079,51 @@ static void ggml_compute_forward_flash_attn_ext_f16( const int ith = params->ith; const int nth = params->nth; +#if defined(GGML_IQK_FLASH_ATTENTION) && defined(__x86_64__) + // opencoti-hook: f5-opt-cpufa — see docs/features/advanced_kv.md + // x86_64-only: the iqk engine defines its symbols under IQK_IMPLEMENT + // (__AVX2__ || __ARM_FEATURE_DOTPROD). Our per-object flags give the iqk TUs + // -Xx86_64-mavx2, so the impl exists on the cosmocc x86_64 slice but NOT the + // aarch64 slice (cosmocc's ARM baseline has no dotprod). Gating the call on + // __x86_64__ keeps it where the symbol lives; ARM falls through to mainline FA. + // Vendored ik_llama.cpp CPU Flash-Attention engine. When enabled at runtime + // and the (type, shape) case is supported, it computes the whole op into + // dst->data and returns true; on false (unsupported case, or alibi via + // max_bias>0 — "noalibi") we fall through to the mainline path below, + // unchanged. Arg mapping mirrors ik_llama's iqk_flash_attn_noalibi call site + // (their ggml.c); llamafile renames params->shared -> params->threadpool. + // n_swa = op_params[4], which is 0 in llamafile (mainline never sets it), so + // the mask carries any sliding-window structure. The wdata scratch is sized + // by the matching iqk_fa_work_buffer_size() bump in ggml-cpu.c's graph plan. + if (ggml_iqk_flash_attn_enabled()) { + const ggml_tensor * iqk_mask = dst->src[3]; + const ggml_tensor * iqk_sinks = dst->src[4]; + float iqk_scale, iqk_max_bias, iqk_softcap; + memcpy(&iqk_scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&iqk_max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&iqk_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (iqk_softcap != 0.0f) { + iqk_scale /= iqk_softcap; + } + if (iqk_flash_attn_noalibi( + q->type, iqk_mask ? iqk_mask->type : GGML_TYPE_F16, iqk_max_bias, + (int) neq3, (int) neq2, (long) nbq3, (long) nbq2, + (int) nek3, (int) nek2, (long) nbk3, (long) nbk2, + (int) nev3, (int) nev2, (long) nbv3, (long) nbv2, + (int) ne2, (int) ne1, (long) nb1, + k->type, v->type, (int) DK, (int) DV, + (int) neq1, (int) nek1, + (int) nbq1, (int) nbk1, (int) nbv1, iqk_mask ? (int) iqk_mask->nb[1] : 0, + q->data, k->data, v->data, + iqk_mask ? iqk_mask->data : NULL, iqk_sinks ? iqk_sinks->data : NULL, + iqk_scale, iqk_softcap, (float *) dst->data, + params->wdata, (barrier_t) ggml_barrier, (void *) params->threadpool, + ith, nth, dst->op_params[4])) { + return; + } + } +#endif + // When use_ref is set, force the vec-only reference implementation (no tiling, no KV-chunking) const bool use_ref = params->use_ref; diff --git a/llama.cpp/ggml/src/ggml-iqk-flash-attn.cpp b/llama.cpp/ggml/src/ggml-iqk-flash-attn.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/ggml-iqk-flash-attn.cpp @@ -0,0 +1,19 @@ +// opencoti F5-opt W2 (#290) — ik_llama CPU Flash-Attention enable switch impl. +// See ggml/include/ggml-iqk-flash-attn.h. Pure ggml; no llama/common include, +// so ggml-cpu/ops.cpp can read the toggle inside the FA dispatch hook. + +#include "ggml-iqk-flash-attn.h" + +#include + +namespace { +std::atomic g_iqk_fa_enabled{false}; +} // namespace + +void ggml_iqk_flash_attn_set_enabled(bool enabled) { + g_iqk_fa_enabled.store(enabled, std::memory_order_relaxed); +} + +bool ggml_iqk_flash_attn_enabled(void) { + return g_iqk_fa_enabled.load(std::memory_order_relaxed); +} diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_128_128.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_128_128.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_128_128.cpp @@ -0,0 +1,45 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +IQK_FA_CASE(iqk_fa_128_128) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + stride_q /= sizeof(float); // q stride as float + auto ck = (const char *)k; + auto cv = (const char *)v; + auto cm = (const char *)mask; + +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + if (type_v != GGML_TYPE_BF16) return false; // we do not support mixing bf16 k-cache with other types + if (nk%64 == 0) { + iqk_flash_helper_T<128, 128, 64>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } + iqk_flash_helper_T<128, 128, 32>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + + if (nk%128 == 0) { + return iqk_flash_helper_T<128, 128, 128>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + if (nk%64 == 0) { + return iqk_flash_helper_T<128, 128, 64>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + + return iqk_flash_helper_T<128, 128, 32>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_192_128.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_192_128.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_192_128.cpp @@ -0,0 +1,45 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +IQK_FA_CASE(iqk_fa_192_128) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + stride_q /= sizeof(float); // q stride as float + auto ck = (const char *)k; + auto cv = (const char *)v; + auto cm = (const char *)mask; + +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + if (type_v != GGML_TYPE_BF16) return false; // we do not support mixing bf16 k-cache with other types + if (nk%64 == 0) { + iqk_flash_helper_T<192, 128, 64>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } + iqk_flash_helper_T<192, 128, 32>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + + if (nk%128 == 0) { + return iqk_flash_helper_T<192, 128, 128>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + if (nk%64 == 0) { + return iqk_flash_helper_T<192, 128, 64>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + + return iqk_flash_helper_T<192, 128, 32>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_192_192.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_192_192.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_192_192.cpp @@ -0,0 +1,45 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +IQK_FA_CASE(iqk_fa_192_192) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + stride_q /= sizeof(float); // q stride as float + auto ck = (const char *)k; + auto cv = (const char *)v; + auto cm = (const char *)mask; + +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + if (type_v != GGML_TYPE_BF16) return false; // we do not support mixing bf16 k-cache with other types + if (nk%64 == 0) { + iqk_flash_helper_T<192, 192, 64>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } + iqk_flash_helper_T<192, 192, 32>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + + if (nk%128 == 0) { + return iqk_flash_helper_T<192, 192, 128>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + if (nk%64 == 0) { + return iqk_flash_helper_T<192, 192, 64>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + + return iqk_flash_helper_T<192, 192, 32>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_256_256.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_256_256.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_256_256.cpp @@ -0,0 +1,45 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +IQK_FA_CASE(iqk_fa_256_256) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + stride_q /= sizeof(float); // q stride as float + auto ck = (const char *)k; + auto cv = (const char *)v; + auto cm = (const char *)mask; + +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + if (type_v != GGML_TYPE_BF16) return false; // we do not support mixing bf16 k-cache with other types + if (nk%64 == 0) { + iqk_flash_helper_T<256, 256, 64>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } + iqk_flash_helper_T<256, 256, 32>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + + if (nk%128 == 0) { + return iqk_flash_helper_T<256, 256, 128>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + if (nk%64 == 0) { + return iqk_flash_helper_T<256, 256, 64>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + + return iqk_flash_helper_T<256, 256, 32>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_320_256.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_320_256.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_320_256.cpp @@ -0,0 +1,142 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +namespace { + +template +inline void iqk_deepseek_helper(KHelper& kh, VHelper& vh, + int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, + const float * q, const char * mask, float scale, float softcap, float * qkv, + const float * sinkf, float * M, float * S) { + auto update = [&nq1, &mask, &q, &qkv, &M, &S, stride_q, stride_m, stride_qkv] (int n) { + nq1 -= n; + if (nq1 == 0) return true; + q += n*stride_q; + mask += n*stride_m; + qkv += n*stride_qkv; + if (M && S) { M += n; S += n; } + return false; + }; + if (nq1 >= 16) { + int n_step = nq1/16; + FlashAttn<320, 256, 16, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 16*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(16*n_step)) return; + } + if (nq1 >= 8) { + int n_step = nq1/8; + FlashAttn<320, 256, 8, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 8*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(8*n_step)) return; + } + if (nq1 >= 4) { + int n_step = nq1/4; + FlashAttn<320, 256, 4, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 4*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(4*n_step)) return; + } + if (nq1 == 3) { + FlashAttn<320, 256, 3, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 3, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } + else if (nq1 == 2) { + FlashAttn<320, 256, 2, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 2, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } else { + FlashAttn<320, 256, 1, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 1, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } +} + +template +inline bool iqk_deepseek_helper(ggml_type type_k, + int nq1, int nk1, int stride_q, int stride_k, int stride_v, int stride_m, int stride_qkv, + const float * q, const char * k, const char * v, const char * mask, + float scale, float softcap, float * qkv, const float * sinkf, float * M, float * S) { + if (type_k == GGML_TYPE_Q8_0) { + HelperQ80 kh((const char *)k, stride_k); + HelperQ80 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q8_0_R8) { + HelperQ80R8<320> kh((const char *)k, stride_k); + HelperQ80 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q6_0) { + HelperQ60 kh((const char *)k, stride_k); + HelperQ60 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#if GGML_IQK_FA_ALL_QUANTS + if (type_k == GGML_TYPE_Q8_KV) { + HelperQ8KV<320> kh((const char *)k, stride_k); + HelperQ8KV<256> vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q4_0) { + HelperQ40 kh((const char *)k, stride_k); + HelperQ40 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q4_1) { + HelperQ41 kh((const char *)k, stride_k); + HelperQ41 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_IQ4_NL) { + HelperIQ4nl kh((const char *)k, stride_k); + HelperIQ4nl vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + if (type_k == GGML_TYPE_F16) { + HelperF16 kh((const char *)k, stride_k); + HelperF16 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + HelperBF16<320, step_k> kh((const char *)k, stride_k); + HelperBF16<256, step_k> vh((const char *)v, stride_v); + if (nq1 % 8 == 0) { + FlashAttnBF16<320, 256, 8, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + } else { + FlashAttnBF16<320, 256, 1, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + } + return true; + } +#endif + return false; +} + +} + +IQK_FA_CASE(iqk_fa_320_256) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + if (!(type_k == type_v || (type_k == GGML_TYPE_Q8_0_R8 && type_v == GGML_TYPE_Q8_0))) { + return false; + } + stride_q /= sizeof(float); // q stride as float + return iqk_deepseek_helper<32>(type_k, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, (const char *)k, (const char *)v, (const char *)mask, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_512_512.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_512_512.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_512_512.cpp @@ -0,0 +1,45 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +IQK_FA_CASE(iqk_fa_512_512) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + stride_q /= sizeof(float); // q stride as float + auto ck = (const char *)k; + auto cv = (const char *)v; + auto cm = (const char *)mask; + +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + if (type_v != GGML_TYPE_BF16) return false; // we do not support mixing bf16 k-cache with other types + if (nk%64 == 0) { + iqk_flash_helper_T<512, 512, 64>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } + iqk_flash_helper_T<512, 512, 32>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + + if (nk%128 == 0) { + return iqk_flash_helper_T<512, 512, 128>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + if (nk%64 == 0) { + return iqk_flash_helper_T<512, 512, 64>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + + return iqk_flash_helper_T<512, 512, 32>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_576_512.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_576_512.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_576_512.cpp @@ -0,0 +1,142 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +namespace { + +template +inline void iqk_deepseek_helper(KHelper& kh, VHelper& vh, + int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, + const float * q, const char * mask, float scale, float softcap, float * qkv, + const float * sinkf, float * M, float * S) { + auto update = [&nq1, &mask, &q, &qkv, &M, &S, stride_q, stride_m, stride_qkv] (int n) { + nq1 -= n; + if (nq1 == 0) return true; + q += n*stride_q; + mask += n*stride_m; + qkv += n*stride_qkv; + if (M && S) { M += n; S += n; } + return false; + }; + if (nq1 >= 16) { + int n_step = nq1/16; + FlashAttn<576, 512, 16, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 16*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(16*n_step)) return; + } + if (nq1 >= 8) { + int n_step = nq1/8; + FlashAttn<576, 512, 8, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 8*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(8*n_step)) return; + } + if (nq1 >= 4) { + int n_step = nq1/4; + FlashAttn<576, 512, 4, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 4*n_step, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + if (update(4*n_step)) return; + } + if (nq1 == 3) { + FlashAttn<576, 512, 3, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 3, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } + else if (nq1 == 2) { + FlashAttn<576, 512, 2, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 2, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } else { + FlashAttn<576, 512, 1, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, 1, nk1, stride_q, stride_m, stride_qkv, q, mask, qkv, M, S); + } +} + +template +inline bool iqk_deepseek_helper(ggml_type type_k, + int nq1, int nk1, int stride_q, int stride_k, int stride_v, int stride_m, int stride_qkv, + const float * q, const char * k, const char * v, const char * mask, + float scale, float softcap, float * qkv, const float * sinkf, float * M, float * S) { + if (type_k == GGML_TYPE_Q8_0) { + HelperQ80 kh((const char *)k, stride_k); + HelperQ80 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q8_0_R8) { + HelperQ80R8<576> kh((const char *)k, stride_k); + HelperQ80 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q6_0) { + HelperQ60 kh((const char *)k, stride_k); + HelperQ60 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#if GGML_IQK_FA_ALL_QUANTS + if (type_k == GGML_TYPE_Q8_KV) { + HelperQ8KV<576> kh((const char *)k, stride_k); + HelperQ8KV<512> vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q4_0) { + HelperQ40 kh((const char *)k, stride_k); + HelperQ40 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_Q4_1) { + HelperQ41 kh((const char *)k, stride_k); + HelperQ41 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } + if (type_k == GGML_TYPE_IQ4_NL) { + HelperIQ4nl kh((const char *)k, stride_k); + HelperIQ4nl vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + if (type_k == GGML_TYPE_F16) { + HelperF16 kh((const char *)k, stride_k); + HelperF16 vh((const char *)v, stride_v); + iqk_deepseek_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + return true; + } +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + HelperBF16<576, step_k> kh((const char *)k, stride_k); + HelperBF16<512, step_k> vh((const char *)v, stride_v); + if (nq1 % 8 == 0) { + FlashAttnBF16<576, 512, 8, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + } else { + FlashAttnBF16<576, 512, 1, step_k> fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + } + return true; + } +#endif + return false; +} + +} + +IQK_FA_CASE(iqk_fa_576_512) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + if (!(type_k == type_v || (type_k == GGML_TYPE_Q8_0_R8 && type_v == GGML_TYPE_Q8_0))) { + return false; + } + stride_q /= sizeof(float); // q stride as float + return iqk_deepseek_helper<32>(type_k, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, (const char *)k, (const char *)v, (const char *)mask, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_64_64.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_64_64.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_64_64.cpp @@ -0,0 +1,45 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +IQK_FA_CASE(iqk_fa_64_64) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + stride_q /= sizeof(float); // q stride as float + auto ck = (const char *)k; + auto cv = (const char *)v; + auto cm = (const char *)mask; + +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + if (type_v != GGML_TYPE_BF16) return false; // we do not support mixing bf16 k-cache with other types + if (nk%64 == 0) { + iqk_flash_helper_T<64, 64, 64>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } + iqk_flash_helper_T<64, 64, 32>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + + if (nk%128 == 0) { + return iqk_flash_helper_T<64, 64, 128>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + if (nk%64 == 0) { + return iqk_flash_helper_T<64, 64, 64>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + + return iqk_flash_helper_T<64, 64, 32>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_96_96.cpp b/llama.cpp/ggml/src/iqk/fa/iqk_fa_96_96.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_96_96.cpp @@ -0,0 +1,45 @@ +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include "iqk/fa/iqk_fa_templates.h" + +IQK_FA_CASE(iqk_fa_96_96) { + + auto type_k = ggml_type(int_type_k); + auto type_v = ggml_type(int_type_v); + + stride_q /= sizeof(float); // q stride as float + auto ck = (const char *)k; + auto cv = (const char *)v; + auto cm = (const char *)mask; + +#ifdef __AVX512BF16__ + if (type_k == GGML_TYPE_BF16) { + if (type_v != GGML_TYPE_BF16) return false; // we do not support mixing bf16 k-cache with other types + if (nk%64 == 0) { + iqk_flash_helper_T<96, 96, 64>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } + iqk_flash_helper_T<96, 96, 32>(nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + return true; + } +#endif + + if (nk%128 == 0) { + return iqk_flash_helper_T<96, 96, 128>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + if (nk%64 == 0) { + return iqk_flash_helper_T<96, 96, 64>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + } + + return iqk_flash_helper_T<96, 96, 32>(type_k, type_v, nq, nk, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, ck, cv, cm, scale, softcap, qkv, sinkf, M, S); + +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/fa/iqk_fa_templates.h b/llama.cpp/ggml/src/iqk/fa/iqk_fa_templates.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/fa/iqk_fa_templates.h @@ -0,0 +1,2258 @@ +// -*- mode:c++;indent-tabs-mode:nil;c-basic-offset:4;coding:utf-8 -*- +// vi: set et ft=cpp fenc=utf-8 :vi +// +// +// Copyright (C) 2024 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include "iqk/iqk_config.h" + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include +#include +#include + +#include "ggml-impl.h" +#include "ggml-quants.h" +#include "iqk/iqk_quantize.h" +#include "iqk/iqk_gemm_floats.h" +#include "iqk/iqk_gemm_kquants.h" +#include "iqk/iqk_gemm_legacy_quants.h" +#include "iqk/iqk_utils.h" + +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" + +// clang-format off + +namespace { + +// Compute effective K boundary by scanning ALL query rows (union-of-masks). +// The original early-termination scanned only the last row, which is correct +// for single-slot (all rows have the same mask) but wrong for multi-slot +// parallel (--parallel N>1) where different slots have different sequence +// lengths and therefore different mask patterns. +// Returns the number of K elements to process (multiple of k_step). +inline int mask_effective_nk1(const char * mask, int n_rows, int stride_m, int nk1, int k_step) { + int ik_max = 0; + for (int j = 0; j < n_rows; ++j) { + auto Mc = (const uint16_t *)(mask + j * stride_m); + int ik = nk1 - k_step; + for (; ik >= 0 && Mc[ik] != 0; ik -= k_step); + ik += k_step; + if (ik > ik_max) ik_max = ik; + } + return ik_max; +} + +struct BaseHelper { + BaseHelper(const char * data, int stride) : data(data), block(data), stride(stride) {} + + //inline void set_block(int k1) { block = data + k1*k_step*stride; } + inline void reset_block() { block = data; } + inline void next_block(int step) { block += step*stride; } + inline const char * lblock(int l1) const { return block + l1*stride; } + + const char * data; + const char * block; + int stride; + +}; + +struct F16 { +#ifdef __AVX512F__ + using Data = __m512; + constexpr static int block_size = 16; + constexpr static int num_registers = 32; + constexpr static int q_step = 8; + static inline Data zero() { return _mm512_setzero_ps(); } + static inline Data load(const char * ptr, int i) { return _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)ptr + i)); } + static inline Data set1(float val) { return _mm512_set1_ps(val); } + static inline Data mul(Data v1, Data v2) { return _mm512_mul_ps(v1, v2); } + static inline Data sub(Data v1, Data v2) { return _mm512_sub_ps(v1, v2); } + static inline Data load(const float * ptr) { return _mm512_loadu_ps(ptr); } + static inline void store(float * ptr, Data data) { _mm512_storeu_ps(ptr, data); } + static inline Data fmadd(Data prev, Data v1, Data v2) { return _mm512_fmadd_ps(v1, v2, prev); } + static inline float reduce_max(Data data) { return _mm512_reduce_max_ps(data); } + static inline float reduce_add(Data data) { return _mm512_reduce_add_ps(data); } + static inline Data max(Data v1, Data v2) { return _mm512_max_ps(v1, v2); } + static inline Data add(Data v1, Data v2) { return _mm512_add_ps(v1, v2); } + static inline Data set4(const float * ptr) { + auto v128 = _mm_loadu_ps(ptr); + auto v256 = _mm256_set_m128(v128, v128); + return _mm512_insertf32x8(_mm512_castps256_ps512(v256), v256, 1); + } + static inline void set4(const float * ptr, Data * vs) { + auto v = set4(ptr); + vs[0] = _mm512_shuffle_ps(v, v, 0x00); + vs[1] = _mm512_shuffle_ps(v, v, 0x55); + vs[2] = _mm512_shuffle_ps(v, v, 0xaa); + vs[3] = _mm512_shuffle_ps(v, v, 0xff); + } + static inline Data fmadd_lane0(Data prev, Data v1, Data v2) { return _mm512_fmadd_ps(v1, _mm512_shuffle_ps(v2, v2, 0x00), prev); } + static inline Data fmadd_lane1(Data prev, Data v1, Data v2) { return _mm512_fmadd_ps(v1, _mm512_shuffle_ps(v2, v2, 0x55), prev); } + static inline Data fmadd_lane2(Data prev, Data v1, Data v2) { return _mm512_fmadd_ps(v1, _mm512_shuffle_ps(v2, v2, 0xaa), prev); } + static inline Data fmadd_lane3(Data prev, Data v1, Data v2) { return _mm512_fmadd_ps(v1, _mm512_shuffle_ps(v2, v2, 0xff), prev); } +#elif defined __AVX2__ + using Data = __m256; + constexpr static int block_size = 8; + constexpr static int num_registers = 16; + constexpr static int q_step = 8; + static inline Data zero() { return _mm256_setzero_ps(); } + static inline Data load(const char * ptr, int i) { return _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)ptr + i)); } + static inline Data set1(float val) { return _mm256_set1_ps(val); } + static inline Data mul(Data v1, Data v2) { return _mm256_mul_ps(v1, v2); } + static inline Data load(const float * ptr) { return _mm256_loadu_ps(ptr); } + static inline Data sub(Data v1, Data v2) { return _mm256_sub_ps(v1, v2); } + static inline void store(float * ptr, Data data) { _mm256_storeu_ps(ptr, data); } + static inline Data fmadd(Data prev, Data v1, Data v2) { return _mm256_fmadd_ps(v1, v2, prev); } + static inline float reduce_max(Data data) { return hmax_float_8(data); } + static inline float reduce_add(Data data) { return hsum_float_8(data); } + static inline Data max(Data v1, Data v2) { return _mm256_max_ps(v1, v2); } + static inline Data add(Data v1, Data v2) { return _mm256_add_ps(v1, v2); } + static inline Data set4(const float * ptr) { + auto v128 = _mm_loadu_ps(ptr); + return _mm256_set_m128(v128, v128); + } + static inline void set4(const float * ptr, Data * vs) { + auto v = set4(ptr); + vs[0] = _mm256_shuffle_ps(v, v, 0x00); + vs[1] = _mm256_shuffle_ps(v, v, 0x55); + vs[2] = _mm256_shuffle_ps(v, v, 0xaa); + vs[3] = _mm256_shuffle_ps(v, v, 0xff); + } + static inline Data fmadd_lane0(Data prev, Data v1, Data v2) { return _mm256_fmadd_ps(v1, _mm256_shuffle_ps(v2, v2, 0x00), prev); } + static inline Data fmadd_lane1(Data prev, Data v1, Data v2) { return _mm256_fmadd_ps(v1, _mm256_shuffle_ps(v2, v2, 0x55), prev); } + static inline Data fmadd_lane2(Data prev, Data v1, Data v2) { return _mm256_fmadd_ps(v1, _mm256_shuffle_ps(v2, v2, 0xaa), prev); } + static inline Data fmadd_lane3(Data prev, Data v1, Data v2) { return _mm256_fmadd_ps(v1, _mm256_shuffle_ps(v2, v2, 0xff), prev); } +#else + using Data = float16x8_t; + constexpr static int block_size = 8; + //constexpr static int num_registers = 32; + //constexpr static int q_step = 8; + static inline Data zero() { return vdupq_n_f16(0); } + static inline Data load(const char * ptr, int i) { return vld1q_f16((const float16_t *)ptr + block_size*i); } + static inline Data load(const float16_t * ptr, int i) { return vld1q_f16(ptr + block_size*i); } + static inline Data load(const float16_t * ptr) { return vld1q_f16(ptr); } + static inline Data load(const float * ptr) { + auto val1 = vld1q_f32(ptr); + auto val2 = vld1q_f32(ptr+4); + return vcombine_f16(vcvt_f16_f32(val1), vcvt_f16_f32(val2)); + } + static inline Data set1(float val) { return vdupq_n_f16(val); } + static inline Data mul(Data v1, Data v2) { return vmulq_f16(v1, v2); } + static inline Data sub(Data v1, Data v2) { return vsubq_f16(v1, v2); } + static inline void store(float * ptr, Data data) { + vst1q_f32(ptr+0, vcvt_f32_f16(vget_low_f16(data))); + vst1q_f32(ptr+4, vcvt_f32_f16(vget_high_f16(data))); + } + static inline void store(float16_t * ptr, Data data) { vst1q_f16(ptr, data); } + static inline void store(float * ptr, float32x4_t data) { vst1q_f32(ptr, data); } + static inline Data fmadd(Data prev, Data v1, Data v2) { return vfmaq_f16(prev, v1, v2); } + static inline float reduce_max(Data data) { return vmaxvq_f16(data); } + static inline float reduce_add(Data data) { + auto sum = vadd_f16(vget_low_f16(data), vget_high_f16(data)); + return vaddvq_f32(vcvt_f32_f16(sum)); + } + static inline Data max(Data v1, Data v2) { return vmaxq_f16(v1, v2); } + static inline Data add(Data v1, Data v2) { return vaddq_f16(v1, v2); } + static inline float16x4_t set4(const float * ptr) { + auto val32 = vld1q_f32(ptr); + return vcvt_f16_f32(val32); + } + static inline Data fmadd_lane0(Data prev, Data v1, float16x4_t v2) { return vfmaq_lane_f16(prev, v1, v2, 0); } + static inline Data fmadd_lane1(Data prev, Data v1, float16x4_t v2) { return vfmaq_lane_f16(prev, v1, v2, 1); } + static inline Data fmadd_lane2(Data prev, Data v1, float16x4_t v2) { return vfmaq_lane_f16(prev, v1, v2, 2); } + static inline Data fmadd_lane3(Data prev, Data v1, float16x4_t v2) { return vfmaq_lane_f16(prev, v1, v2, 3); } +#endif + template static inline float reduce_max(const Data * data) { + return reduce_T(data); + } + template static inline float reduce_add(const Data * data) { + return reduce_T(data); + } + template + static float reduce_T(const Data * data) { + float result; + if constexpr (k_step/block_size == 1) { + result = Op(data[0]); + } + else if constexpr (k_step/block_size == 2) { + result = Op(Op_combine(data[0], data[1])); + } + else { + auto vmax = Op_combine(data[0], data[1]); + for (int l = 2; l < k_step/block_size; ++l) vmax = Op_combine(vmax, data[l]); + result = Op(vmax); + } + return result; + } +}; + +struct HelperF16 final : public BaseHelper { + using Base = BaseHelper; + HelperF16(const char * data, int stride) : Base(data, stride) {} + + inline void load(int l1, int i, F16::Data& v1, F16::Data& v2) const { + //auto dr = (const ggml_half *)Base::lblock(l1); + auto dr = Base::lblock(l1); + v1 = F16::load(dr, i + 0); + v2 = F16::load(dr, i + 1); + } +}; + +template struct block_q8_KV { + float d; + int s; + int8_t qs[D]; +}; + +template +struct HelperQ8KV final : public BaseHelper { + using Base = BaseHelper; + using block_q8 = block_q8_KV; + constexpr static ggml_type type = GGML_TYPE_Q8_KV; + constexpr static int block_size_q = D; + HelperQ8KV(const char * data, int stride) : Base(data, stride) {} + + // Needed for v * softmax(k * q) + inline void load(int l1, int i, F16::Data& v1, F16::Data& v2) const { + auto q8 = (const block_q8_KV *)Base::lblock(l1); +#ifdef __aarch64__ + auto vd = F16::set1(q8->d); + auto qs = vld1_s8_x2(q8->qs + 8*i); + v1 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(qs.val[0]))); + v2 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(qs.val[1]))); +#else + auto vd = F16::set1(q8->d); +#ifdef __AVX512F__ + v1 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)q8->qs+i+0)))); + v2 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)q8->qs+i+1)))); +#else + v1 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *)(q8->qs+8*i+0))))); + v2 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *)(q8->qs+8*i+8))))); +#endif +#endif + } +}; + +struct HelperQ80 final : public BaseHelper { + using Base = BaseHelper; + constexpr static ggml_type type = GGML_TYPE_Q8_0; +//#ifdef HAVE_FANCY_SIMD +#ifdef __AVX2__ + using block_q8 = block_q8_2; + constexpr static int block_size_q = QK8_2; +#else + using block_q8 = block_q8_0; + constexpr static int block_size_q = QK8_0; +#endif + HelperQ80(const char * data, int stride) : Base(data, stride) {} + + // Needed for v * softmax(k * q) + inline void load(int l1, int i, F16::Data& v1, F16::Data& v2) const { + int j = F16::block_size*i; + auto dl = (const block_q8_0 *)Base::lblock(l1) + j/QK8_0; +#ifdef __aarch64__ + auto vd = F16::set1(GGML_FP16_TO_FP32(dl->d)); + int ii = j%QK8_0; + auto qs = vld1_s8_x2(dl->qs + ii); + v1 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(qs.val[0]))); + v2 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(qs.val[1]))); +#else + auto vd = F16::set1(GGML_FP16_TO_FP32(dl->d)); +#ifdef __AVX512F__ + v1 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)dl->qs+0)))); + v2 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)dl->qs+1)))); +#else + int ii = j%QK8_0; + v1 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *)(dl->qs+ii+0))))); + v2 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *)(dl->qs+ii+8))))); +#endif +#endif + } + + template + static inline void convert(int nq, int stride_q, const float * q, block_q8_0 * y) { + for (int i = 0; i < nq; ++i) { + quantize_row_q8_0_x4(q, y, D); + q += stride_q; + y += D/QK8_0; + } + } + + template + static inline void convert(int nq, int stride_q, const float * q, block_q8_1 * y) { + for (int i = 0; i < nq; ++i) { + quantize_row_q8_1_x4(q, y, D); + q += stride_q; + y += D/QK8_1; + } + } + + template + static inline void convert(int nq, int stride_q, const float * q, block_q8_2 * y) { + for (int i = 0; i < nq; ++i) { + quantize_row_q8_2_x4(q, y, D); + q += stride_q; + y += D/QK8_2; + } + } + + template + static inline void convert(int nq, int stride_q, const float * q, block_q8_KV * y) { + for (int i = 0; i < nq; ++i) { + quantize_row_q8_KV(q, y, D); + q += stride_q; + ++y; + } + } +}; + +template +struct HelperQ80R8 : public BaseHelper { + using Base = BaseHelper; + constexpr static ggml_type type = GGML_TYPE_Q8_0_R8; +#ifdef __AVX2__ + constexpr static int block_size_q = QK8_2; + using block_q8 = block_q8_2; +#else + constexpr static int block_size_q = QK8_0; + using block_q8 = block_q8_0; +#endif + HelperQ80R8(const char * data, int stride) : Base(data, stride) {} + HelperQ80R8(int nk, const HelperQ80& q8) : Base(q8.data, q8.stride) { + r4 = repack(nk, q8); + Base::data = (const char *)r4.data(); + Base::stride = (D/QK8_0)*sizeof(block_q8_0); + } + + static void repack(int nk, const char * q8_data, int q8_stride, block_q8_0_r8 * y) { + constexpr int nblock = D/QK8_0; + const block_q8_0 * x8[8]; +#ifdef __ARM_NEON + int8x16x2_t m0, m1, m2, m3; +#endif + for (int row = 0; row < nk; row += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q8_0 *)(q8_data + (row + k)*q8_stride); + for (int ib = 0; ib < nblock; ++ib) { + for (int k = 0; k < 8; ++k) y[ib].d[k] = x8[k][ib].d; +#ifdef __AVX2__ + auto m0 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[4][ib].qs), _mm_loadu_si128((const __m128i *)x8[0][ib].qs)); + auto m1 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[5][ib].qs), _mm_loadu_si128((const __m128i *)x8[1][ib].qs)); + auto m2 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[6][ib].qs), _mm_loadu_si128((const __m128i *)x8[2][ib].qs)); + auto m3 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[7][ib].qs), _mm_loadu_si128((const __m128i *)x8[3][ib].qs)); + auto t0 = _mm256_unpacklo_epi32(m0, m1); + auto t1 = _mm256_unpacklo_epi32(m2, m3); + auto t2 = _mm256_unpackhi_epi32(m0, m1); + auto t3 = _mm256_unpackhi_epi32(m2, m3); + m0 = _mm256_unpacklo_epi64(t0, t1); + m1 = _mm256_unpackhi_epi64(t0, t1); + m2 = _mm256_unpacklo_epi64(t2, t3); + m3 = _mm256_unpackhi_epi64(t2, t3); +//#ifdef HAVE_FANCY_SIMD +// m0 = _mm256_add_epi8(m0, _mm256_set1_epi8(127)); +// m1 = _mm256_add_epi8(m1, _mm256_set1_epi8(127)); +// m2 = _mm256_add_epi8(m2, _mm256_set1_epi8(127)); +// m3 = _mm256_add_epi8(m3, _mm256_set1_epi8(127)); +//#endif + _mm256_storeu_si256((__m256i *)y[ib].qs + 0, m0); + _mm256_storeu_si256((__m256i *)y[ib].qs + 1, m1); + _mm256_storeu_si256((__m256i *)y[ib].qs + 2, m2); + _mm256_storeu_si256((__m256i *)y[ib].qs + 3, m3); + m0 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[4][ib].qs+1), _mm_loadu_si128((const __m128i *)x8[0][ib].qs+1)); + m1 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[5][ib].qs+1), _mm_loadu_si128((const __m128i *)x8[1][ib].qs+1)); + m2 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[6][ib].qs+1), _mm_loadu_si128((const __m128i *)x8[2][ib].qs+1)); + m3 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[7][ib].qs+1), _mm_loadu_si128((const __m128i *)x8[3][ib].qs+1)); + t0 = _mm256_unpacklo_epi32(m0, m1); + t1 = _mm256_unpacklo_epi32(m2, m3); + t2 = _mm256_unpackhi_epi32(m0, m1); + t3 = _mm256_unpackhi_epi32(m2, m3); + m0 = _mm256_unpacklo_epi64(t0, t1); + m1 = _mm256_unpackhi_epi64(t0, t1); + m2 = _mm256_unpacklo_epi64(t2, t3); + m3 = _mm256_unpackhi_epi64(t2, t3); +//#ifdef HAVE_FANCY_SIMD +// m0 = _mm256_add_epi8(m0, _mm256_set1_epi8(127)); +// m1 = _mm256_add_epi8(m1, _mm256_set1_epi8(127)); +// m2 = _mm256_add_epi8(m2, _mm256_set1_epi8(127)); +// m3 = _mm256_add_epi8(m3, _mm256_set1_epi8(127)); +//#endif + _mm256_storeu_si256((__m256i *)y[ib].qs + 4, m0); + _mm256_storeu_si256((__m256i *)y[ib].qs + 5, m1); + _mm256_storeu_si256((__m256i *)y[ib].qs + 6, m2); + _mm256_storeu_si256((__m256i *)y[ib].qs + 7, m3); +#elif defined __ARM_NEON + for (int l = 0; l < 2; ++l) { + m0.val[0] = vld1q_s8(x8[0][ib].qs+16*l); m0.val[1] = vld1q_s8(x8[4][ib].qs+16*l); + m1.val[0] = vld1q_s8(x8[1][ib].qs+16*l); m1.val[1] = vld1q_s8(x8[5][ib].qs+16*l); + m2.val[0] = vld1q_s8(x8[2][ib].qs+16*l); m2.val[1] = vld1q_s8(x8[6][ib].qs+16*l); + m3.val[0] = vld1q_s8(x8[3][ib].qs+16*l); m3.val[1] = vld1q_s8(x8[7][ib].qs+16*l); + auto row01 = vtrnq_s32(vreinterpretq_s32_s8(m0.val[0]), vreinterpretq_s32_s8(m1.val[0])); + auto row23 = vtrnq_s32(vreinterpretq_s32_s8(m2.val[0]), vreinterpretq_s32_s8(m3.val[0])); + m0.val[0] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m1.val[0] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + m2.val[0] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m3.val[0] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + row01 = vtrnq_s32(vreinterpretq_s32_s8(m0.val[1]), vreinterpretq_s32_s8(m1.val[1])); + row23 = vtrnq_s32(vreinterpretq_s32_s8(m2.val[1]), vreinterpretq_s32_s8(m3.val[1])); + m0.val[1] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m1.val[1] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + m2.val[1] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m3.val[1] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + vst1q_s8_x2(y[ib].qs + 0 + 128*l, m0); + vst1q_s8_x2(y[ib].qs + 32 + 128*l, m1); + vst1q_s8_x2(y[ib].qs + 64 + 128*l, m2); + vst1q_s8_x2(y[ib].qs + 96 + 128*l, m3); + } +#else + for (int l = 0; l < 4; ++l) { + for (int k = 0; k < 8; ++k) for (int i = 0; i < 4; ++i) { + y[ib].qs[32*l+4*k+i+ 0] = x8[k][ib].qs[i+4*l+ 0]; + y[ib].qs[32*l+4*k+i+128] = x8[k][ib].qs[i+4*l+16]; + } + } +#endif + } + y += nblock; + } + } + + static std::vector repack(int nk, const HelperQ80& q8) { + static_assert(D%QK8_0 == 0); + GGML_ASSERT(nk%8 == 0); + constexpr int nblock = D/QK8_0; + std::vector result(nblock * nk/8); + auto y = result.data(); + repack(nk, q8.data, q8.stride, y); + return result; + } + + std::vector r4; +}; + +// TODO: unite this with the above +template +struct HelperQ8KVR8 : public BaseHelper { + using Base = BaseHelper; + constexpr static ggml_type type = GGML_TYPE_Q8_KV_R8; + constexpr static int block_size_q = D; + using block_q8 = block_q8_KV; + + struct block_q8_KV_r8 { + float d[8]; + int8_t qs[8*D]; + }; + + HelperQ8KVR8(int nk, const HelperQ8KV& q8) : Base(q8.data, q8.stride) { + r4 = repack(nk, q8); + Base::data = (const char *)r4.data(); + Base::stride = sizeof(block_q8_KV_r8)/8; + } + + static std::vector repack(int nk, const HelperQ8KV& q8) { + static_assert(D%32 == 0); + GGML_ASSERT(nk%8 == 0); + std::vector result(nk/8); + auto y = result.data(); +#ifdef __ARM_NEON + int8x16x2_t m0, m1, m2, m3; +#endif + const int8_t * x8[8]; + for (int ix = 0; ix < nk/8; ++ix) { + for (int k = 0; k < 8; ++k) { + auto dptr = (const float *)(q8.data + (8*ix + k)*q8.stride); + y[ix].d[k] = dptr[0]; + x8[k] = (const int8_t *)(dptr + 2); + } + for (int ib = 0; ib < D/16; ++ib) { +#ifdef __AVX2__ + auto m0 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[4]+ib), _mm_loadu_si128((const __m128i *)x8[0]+ib)); + auto m1 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[5]+ib), _mm_loadu_si128((const __m128i *)x8[1]+ib)); + auto m2 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[6]+ib), _mm_loadu_si128((const __m128i *)x8[2]+ib)); + auto m3 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[7]+ib), _mm_loadu_si128((const __m128i *)x8[3]+ib)); + auto t0 = _mm256_unpacklo_epi32(m0, m1); + auto t1 = _mm256_unpacklo_epi32(m2, m3); + auto t2 = _mm256_unpackhi_epi32(m0, m1); + auto t3 = _mm256_unpackhi_epi32(m2, m3); + m0 = _mm256_unpacklo_epi64(t0, t1); + m1 = _mm256_unpackhi_epi64(t0, t1); + m2 = _mm256_unpacklo_epi64(t2, t3); + m3 = _mm256_unpackhi_epi64(t2, t3); +//#ifdef HAVE_FANCY_SIMD +// m0 = _mm256_add_epi8(m0, _mm256_set1_epi8(127)); +// m1 = _mm256_add_epi8(m1, _mm256_set1_epi8(127)); +// m2 = _mm256_add_epi8(m2, _mm256_set1_epi8(127)); +// m3 = _mm256_add_epi8(m3, _mm256_set1_epi8(127)); +//#endif + _mm256_storeu_si256((__m256i *)y[ix].qs + 4*ib+0, m0); + _mm256_storeu_si256((__m256i *)y[ix].qs + 4*ib+1, m1); + _mm256_storeu_si256((__m256i *)y[ix].qs + 4*ib+2, m2); + _mm256_storeu_si256((__m256i *)y[ix].qs + 4*ib+3, m3); +#elif defined __ARM_NEON + // TODO + m0.val[0] = vld1q_s8(x8[0]+16*ib); m0.val[1] = vld1q_s8(x8[4]+16*ib); + m1.val[0] = vld1q_s8(x8[1]+16*ib); m1.val[1] = vld1q_s8(x8[5]+16*ib); + m2.val[0] = vld1q_s8(x8[2]+16*ib); m2.val[1] = vld1q_s8(x8[6]+16*ib); + m3.val[0] = vld1q_s8(x8[3]+16*ib); m3.val[1] = vld1q_s8(x8[7]+16*ib); + auto row01 = vtrnq_s32(vreinterpretq_s32_s8(m0.val[0]), vreinterpretq_s32_s8(m1.val[0])); + auto row23 = vtrnq_s32(vreinterpretq_s32_s8(m2.val[0]), vreinterpretq_s32_s8(m3.val[0])); + m0.val[0] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m1.val[0] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + m2.val[0] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m3.val[0] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + row01 = vtrnq_s32(vreinterpretq_s32_s8(m0.val[1]), vreinterpretq_s32_s8(m1.val[1])); + row23 = vtrnq_s32(vreinterpretq_s32_s8(m2.val[1]), vreinterpretq_s32_s8(m3.val[1])); + m0.val[1] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m1.val[1] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + m2.val[1] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m3.val[1] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + vst1q_s8_x2(y[ix].qs + 0 + 128*ib, m0); + vst1q_s8_x2(y[ix].qs + 32 + 128*ib, m1); + vst1q_s8_x2(y[ix].qs + 64 + 128*ib, m2); + vst1q_s8_x2(y[ix].qs + 96 + 128*ib, m3); +#else + // TODO + for (int l = 0; l < 4; ++l) { + for (int k = 0; k < 8; ++k) for (int i = 0; i < 4; ++i) { + y[ib].qs[32*l+4*k+i+ 0] = x8[k][ib].qs[i+4*l+ 0]; + y[ib].qs[32*l+4*k+i+128] = x8[k][ib].qs[i+4*l+16]; + } + } +#endif + } + } + return result; + } + + std::vector r4; +}; + +struct HelperQ40 final : public BaseHelper { + using Base = BaseHelper; + constexpr static ggml_type type = GGML_TYPE_Q4_0; +#if defined __AVX2__ + using block_q8 = block_q8_2; + constexpr static int block_size_q = QK8_2; +#else + using block_q8 = block_q8_0; + constexpr static int block_size_q = QK8_0; +#endif + HelperQ40(const char * data, int stride) : Base(data, stride) {} + + // Needed for v * softmax(k * q) + inline void load(int l1, int i, F16::Data& v1, F16::Data& v2) const { + int j = F16::block_size*i; + auto dl = (const block_q4_0 *)Base::lblock(l1) + j/QK4_0; +#ifdef __aarch64__ + auto vd = F16::set1(*(const float16_t *)&dl->d); + auto q = vld1q_u8(dl->qs); + q = j%QK4_0 ? vshrq_n_u8(q, 4) : vandq_u8(q, mask); + q = vaddq_s8(q, m8); + v1 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(vget_low_s8(q)))); + v2 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(vget_high_s8(q)))); +#else + auto vd = F16::set1(GGML_FP16_TO_FP32(dl->d)); + auto q = _mm_loadu_si128((const __m128i *)dl->qs); +#ifdef __AVX512F__ + auto ql = _mm_add_epi8(_mm_and_si128(q, mask), m8); + auto qh = _mm_add_epi8(_mm_and_si128(_mm_srli_epi16(q, 4), mask), m8); + v1 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(ql))); + v2 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(qh))); +#else + if (j%QK4_0) q = _mm_srli_epi16(q, 4); + auto q16 = _mm256_cvtepi8_epi16(_mm_add_epi8(_mm_and_si128(q, mask), m8)); + v1 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16)))); + v2 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16, 1)))); +#endif +#endif + } + +#ifdef __AVX2__ + const __m128i mask = _mm_set1_epi8(0xf); + const __m128i m8 = _mm_set1_epi8(-8); +#else + const uint8x16_t mask = vdupq_n_u8(0xf); + const int8x16_t m8 = vdupq_n_s8(-8); +#endif +}; + +struct HelperQ41 final : public BaseHelper { + using Base = BaseHelper; + using block_q8 = block_q8_2; + constexpr static ggml_type type = GGML_TYPE_Q4_1; + constexpr static int block_size_q = QK8_2; + HelperQ41(const char * data, int stride) : Base(data, stride) {} + + // Needed for v * softmax(k * q) + inline void load(int l1, int i, F16::Data& v1, F16::Data& v2) const { + int j = F16::block_size*i; + auto dl = (const block_q4_1 *)Base::lblock(l1) + j/QK4_1; +#ifdef __aarch64__ + auto vd = F16::set1(*(const float16_t *)&dl->d); + auto vm = F16::set1(*(const float16_t *)&dl->m); + auto q = vld1q_u8(dl->qs); + q = (j%QK4_1) ? vshrq_n_u8(q, 4) : vandq_u8(q, mask); + v1 = vfmaq_f16(vm, vd, vcvtq_f16_u16(vmovl_u8(vget_low_u8(q)))); + v2 = vfmaq_f16(vm, vd, vcvtq_f16_u16(vmovl_u8(vget_high_u8(q)))); +#else + auto vd = F16::set1(GGML_FP16_TO_FP32(dl->d)); + auto vm = F16::set1(GGML_FP16_TO_FP32(dl->m)); + auto q = _mm_loadu_si128((const __m128i *)dl->qs); +#ifdef __AVX512F__ + auto ql = _mm_and_si128(q, mask); + auto qh = _mm_and_si128(_mm_srli_epi16(q, 4), mask); + v1 = _mm512_fmadd_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(ql)), vm); + v2 = _mm512_fmadd_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(qh)), vm); +#else + if (j%QK4_1) q = _mm_srli_epi16(q, 4); + auto q16 = _mm256_cvtepi8_epi16(_mm_and_si128(q, mask)); + v1 = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16))), vm); + v2 = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16, 1))), vm); +#endif +#endif + } + +#ifdef __aarch64__ + const uint8x16_t mask = vdupq_n_u8(0xf); +#else + const __m128i mask = _mm_set1_epi8(0xf); +#endif +}; + +struct HelperIQ4nl final : public BaseHelper { + using Base = BaseHelper; + constexpr static ggml_type type = GGML_TYPE_IQ4_NL; +#ifdef __aarch64__ + using block_q8 = block_q8_0; + HelperIQ4nl(const char * data, int stride) : Base(data, stride), values(vld1q_s8(iq4k_values)) {} + constexpr static int block_size_q = QK8_0; +#else + HelperIQ4nl(const char * data, int stride) : Base(data, stride) {} + using block_q8 = block_q8_2; + constexpr static int block_size_q = QK8_2; +#endif + + // Needed for v * softmax(k * q) + inline void load(int l1, int i, F16::Data& v1, F16::Data& v2) const { + int j = F16::block_size*i; + auto dl = (const block_iq4_nl *)Base::lblock(l1) + j/QK4_0; +#ifdef __aarch64__ + auto vd = F16::set1(*(const float16_t *)&dl->d); + auto q = vld1q_u8(dl->qs); + q = j%QK4_0 ? vshrq_n_u8(q, 4) : vandq_u8(q, mask); + q = vqtbl1q_s8(values, q); + v1 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(vget_low_s8(q)))); + v2 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(vget_high_s8(q)))); +#else + auto vd = F16::set1(GGML_FP16_TO_FP32(dl->d)); + auto q = _mm_loadu_si128((const __m128i *)dl->qs); +#ifdef __AVX512F__ + auto ql = _mm_shuffle_epi8(values, _mm_and_si128(q, mask)); + auto qh = _mm_shuffle_epi8(values, _mm_and_si128(_mm_srli_epi16(q, 4), mask)); + v1 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(ql))); + v2 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(qh))); +#else + if (j%QK4_0) q = _mm_srli_epi16(q, 4); + auto q16 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(values, _mm_and_si128(q, mask))); + v1 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16)))); + v2 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16, 1)))); +#endif +#endif + } + +#ifdef __aarch64__ + const uint8x16_t mask = vdupq_n_u8(0xf); + const int8x16_t values; +#else + const __m128i mask = _mm_set1_epi8(0xf); + const __m128i values = _mm_loadu_si128((const __m128i *)iq4k_values); +#endif +}; + +struct HelperQ60 final : public BaseHelper { + constexpr static ggml_type type = GGML_TYPE_Q6_0; +#ifdef __aarch64__ + using block_q8 = block_q8_0; + constexpr static int block_size_q = QK8_0; +#else + using block_q8 = block_q8_2; + constexpr static int block_size_q = QK8_2; +#endif + using Base = BaseHelper; + HelperQ60(const char * data, int stride) : Base(data, stride) {} + + // Needed for v * softmax(k * q) + inline void load(int l1, int i, F16::Data& v1, F16::Data& v2) const { + int j = F16::block_size*i; + auto dl = (const block_q6_0 *)Base::lblock(l1) + j/QK6_0; +#ifdef __aarch64__ + // TODO + const float16_t * d16 = (const float16_t *)&dl->d; + auto vd = F16::set1(d16[0]); + //auto vd = F16::set1(*(const float16_t *)&dl->d); + auto qh8 = vld1_u8(dl->qh); + auto qh = vcombine_u8(vshl_n_u8(qh8, 4), qh8); + auto qs = vld1q_u8(dl->qs); + qs = j%QK4_0 ? vshrq_n_u8(qs, 4) : vandq_u8(qs, mask_l); + qs = vorrq_u8(qs, vandq_u8(mask_h, j%QK4_0 ? vshrq_n_u8(qh, 2) : qh)); + qs = vaddq_s8(qs, m32); + v1 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(vget_low_s8(qs)))); + v2 = vmulq_f16(vd, vcvtq_f16_s16(vmovl_s8(vget_high_s8(qs)))); +#else + auto vd = F16::set1(GGML_FP16_TO_FP32(dl->d)); + auto bl = _mm_loadu_si128((const __m128i *)dl->qs); + uint64_t aux64; std::memcpy(&aux64, dl->qh, 8); + auto bh = _mm_set_epi64x(aux64, aux64 << 4); +#ifdef __AVX512F__ + auto ql = _mm_add_epi8(_mm_or_si128(_mm_and_si128(bl, mask_l), _mm_and_si128(bh, mask_h)), m32); + auto qh = _mm_add_epi8(_mm_or_si128(_mm_and_si128(_mm_srli_epi16(bl, 4), mask_l), _mm_and_si128(_mm_srli_epi16(bh, 2), mask_h)), m32); + v1 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(ql))); + v2 = _mm512_mul_ps(vd, _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(qh))); +#else + if (j%QK4_0) { + bl = _mm_srli_epi16(bl, 4); + bh = _mm_srli_epi16(bh, 2); + } + auto q16 = _mm256_cvtepi8_epi16(_mm_add_epi8(_mm_or_si128(_mm_and_si128(bl, mask_l), _mm_and_si128(bh, mask_h)), m32)); + v1 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16)))); + v2 = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16, 1)))); +#endif +#endif + } + +#ifdef __AVX2__ + const __m128i mask_l = _mm_set1_epi8(0x0f); + const __m128i mask_h = _mm_set1_epi8(0x30); + const __m128i m32 = _mm_set1_epi8(-32); +#else + const uint8x16_t mask_l = vdupq_n_u8(0x0f); + const uint8x16_t mask_h = vdupq_n_u8(0x30); + const int8x16_t m32 = vdupq_n_s8(-32); +#endif +}; + +template +struct FlashMS { + constexpr static int q_step = q_step_in; + constexpr static int k_step = k_step_in; +// Something goes wrong when storing and manipulating K*Q as fp16. +// It works for some models (e.g., Gemma-2), but not for others (e.g., LLaMA-3.1-8B). +// As I wasn't able to find where we lose precision, let's comment this out +// for now and do the K*Q part in fp32. +//#ifdef __aarch64__ +// using cache_t = float16_t; +//#else +// using cache_t = float; +//#endif + using cache_t = float; + + FlashMS(float scale, float softcap) : vscale(F16::set1(scale)), softcap(softcap), h_inf(GGML_FP32_TO_FP16(-INFINITY)) {} + + inline void init_qstep() { + for (int j = 0; j < q_step; ++j) { + S[j] = 0; M[j] = -INFINITY; + } + } + + inline void update_M(int j, float smax) { + if (smax == -INFINITY) { + std::memset(cache + k_step*j, 0, k_step*sizeof(float)); + need_scaling[j] = M[j] == -INFINITY ? 2 : 0; + return; + } + need_scaling[j] = 0; + if (smax > M[j]) { + if (M[j] > -INFINITY) { + float m = expf(M[j] - smax); + vms[j] = m; + need_scaling[j] = 1; + S[j] *= m; + } else { + need_scaling[j] = 2; + S[j] = 0; + } + M[j] = smax; + } + } + +#ifdef __aarch64__ + inline void update_S(int j, float32x4_t * vk) { + auto vm = vdupq_n_f32(M[j]); + auto vsum = vdupq_n_f32(0); + for (int l = 0; l < k_step/4; ++l) { + vk[l] = v_expf(vsubq_f32(vk[l], vm)); + vsum = vaddq_f32(vsum, vk[l]); + F16::store(cache + k_step*j + 4*l, vk[l]); + } + S[j] += vaddvq_f32(vsum); + } +#else + inline void update_S(int j, F16::Data * vk) { + auto vm = F16::set1(M[j]); + for (int l = 0; l < k_step/F16::block_size; ++l) { + vk[l] = v_expf(F16::sub(vk[l], vm)); + F16::store(cache + k_step*j + F16::block_size*l, vk[l]); + } + S[j] += F16::reduce_add(vk); + } +#endif + +#ifdef __aarch64__ + inline float load_and_scale(int j, float32x4_t * vk) { + float32x4_t vmax = vdupq_n_f32(-INFINITY); + // Something goes wrong when storing and manipulating K*Q as fp16. + // It works for some models (e.g., Gemma-2), but not for others (e.g., LLaMA-3.1-8B). + // As I wasn't able to find where we lose precision, let's comment this out + // for now and do the K*Q part in fp32. + //if (softcap <= 0.0f) { + // for (int l = 0; l < k_step/F16::block_size; ++l) { + // auto val = F16::mul(vscale, F16::load(cache + k_step*j + F16::block_size*l)); + // vk[2*l+0] = vcvt_f32_f16(vget_low_f16(val)); + // vk[2*l+1] = vcvt_f32_f16(vget_high_f16(val)); + // vmax = vmaxq_f32(vmax, vmaxq_f32(vk[2*l+0], vk[2*l+1])); + // } + //} else { + // auto v_softcap = vdupq_n_f32(softcap); + // for (int l = 0; l < k_step/F16::block_size; ++l) { + // auto val = F16::mul(vscale, F16::load(cache + k_step*j + F16::block_size*l)); + // vk[2*l+0] = vcvt_f32_f16(vget_low_f16(val)); + // vk[2*l+1] = vcvt_f32_f16(vget_high_f16(val)); + // vk[2*l+0] = vmulq_f32(v_softcap, v_tanh(vk[2*l+0])); + // vk[2*l+1] = vmulq_f32(v_softcap, v_tanh(vk[2*l+1])); + // vmax = vmaxq_f32(vmax, vmaxq_f32(vk[2*l+0], vk[2*l+1])); + // } + //} + auto vscale32 = vcvt_f32_f16(vget_low_f16(vscale)); + if (softcap <= 0.0f) { + for (int l = 0; l < k_step/4; ++l) { + vk[l] = vmulq_f32(vscale32, vld1q_f32(cache + k_step*j + 4*l)); + vmax = vmaxq_f32(vmax, vk[l]); + } + } else { + auto v_softcap = vdupq_n_f32(softcap); + for (int l = 0; l < k_step/4; ++l) { + vk[l] = vmulq_f32(vscale32, vld1q_f32(cache + k_step*j + 4*l)); + vk[l] = vmulq_f32(v_softcap, v_tanh(vk[l])); + vmax = vmaxq_f32(vmax, vk[l]); + } + } + return vmaxvq_f32(vmax); + } + inline float load_apply_mask_and_scale(int j, float32x4_t * vk, const char * mask) { + auto vzero = vdupq_n_f16(0); + auto vinf = vdupq_n_f32(-INFINITY); + for (int l = 0; l < k_step/8; ++l) { + auto vm = vceqq_f16(vzero, vld1q_f16((const float16_t *)mask + 8*l)); + auto vm1 = vzip1q_u16(vm, vm); + auto vm2 = vzip2q_u16(vm, vm); + auto kq = vld1q_f32_x2(cache + k_step*j + 8*l); + vk[2*l+0] = vreinterpretq_f32_u32(vorrq_u32(vandq_u32(vreinterpretq_u32_f32(kq.val[0]), vm1), + vbicq_u32(vreinterpretq_u32_f32(vinf), vm1))); + vk[2*l+1] = vreinterpretq_f32_u32(vorrq_u32(vandq_u32(vreinterpretq_u32_f32(kq.val[1]), vm2), + vbicq_u32(vreinterpretq_u32_f32(vinf), vm2))); + } + float32x4_t vmax = vdupq_n_f32(-INFINITY); + auto vscale32 = vcvt_f32_f16(vget_low_f16(vscale)); + if (softcap <= 0.0f) { + for (int l = 0; l < k_step/4; ++l) { + vk[l] = vmulq_f32(vscale32, vk[l]); + vmax = vmaxq_f32(vmax, vk[l]); + } + } else { + auto v_softcap = vdupq_n_f32(softcap); + for (int l = 0; l < k_step/4; ++l) { + vk[l] = vmulq_f32(vscale32, vk[l]); + vk[l] = vmulq_f32(v_softcap, v_tanh(vk[l])); + vmax = vmaxq_f32(vmax, vk[l]); + } + } + return vmaxvq_f32(vmax); + } +#else + inline float load_and_scale(int j, F16::Data * vk) { + if (softcap <= 0.0f) { + for (int l = 0; l < k_step/F16::block_size; ++l) vk[l] = F16::mul(vscale, F16::load(cache + k_step*j + F16::block_size*l)); + } else { + auto v_softcap = F16::set1(softcap); + for (int l = 0; l < k_step/F16::block_size; ++l) { + auto val = F16::load(cache + k_step*j + F16::block_size*l); + vk[l] = F16::mul(v_softcap, v_tanh(F16::mul(vscale, val))); + } + } + return F16::reduce_max(vk); + } + static inline __m256 apply_mask(int l, const char * mask, __m256 val, [[maybe_unused]] __m256 vinf) { + return _mm256_add_ps(val, _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)mask+l))); + //auto m128 = _mm_loadu_si128((const __m128i *)mask+l); + //m128 = _mm_cmpeq_epi16(m128, _mm_setzero_si128()); + //auto m256 = _mm256_cvtepi16_epi32(m128); + //auto mf = _mm256_castsi256_ps(_mm256_or_si256(m256, _mm256_slli_epi32(m256, 16))); + //return _mm256_or_ps(_mm256_and_ps(mf, val), _mm256_andnot_ps(mf, vinf)); + } +#ifdef __AVX512F__ + static inline __m512 apply_mask(int l, const char * mask, __m512 val, __m512 vinf) { + auto m256 = _mm256_loadu_si256((const __m256i *)mask+l); + m256 = _mm256_cmpeq_epi16(m256, _mm256_setzero_si256()); + auto m512 = _mm512_cvtepi16_epi32(m256); + auto mf = _mm512_castsi512_ps(_mm512_or_si512(m512, _mm512_slli_epi32(m512, 16))); + return _mm512_or_ps(_mm512_and_ps(mf, val), _mm512_andnot_ps(mf, vinf)); + } +#endif + inline float load_apply_mask_and_scale(int j, F16::Data * vk, const char * mask) { +#ifdef HAVE_FANCY_SIMD + auto vzero = _mm256_set1_epi16(0); + auto vinf = _mm512_set1_ps(-INFINITY); + if (softcap <= 0) { + for (int l = 0; l < k_step/F16::block_size; ++l) { + auto m16 = _mm256_cmpeq_epi16_mask(_mm256_loadu_si256((const __m256i *)mask + l), vzero); + vk[l] = _mm512_mask_mul_ps(vinf, m16, vscale, F16::load(cache + k_step*j + F16::block_size*l)); + } + } else { + auto v_softcap = F16::set1(softcap); + for (int l = 0; l < k_step/F16::block_size; ++l) { + auto m16 = _mm256_cmpeq_epi16_mask(_mm256_loadu_si256((const __m256i *)mask + l), vzero); + vk[l] = _mm512_mask_mul_ps(vinf, m16, v_softcap, v_tanh(F16::mul(vscale, F16::load(cache + k_step*j + F16::block_size*l)))); + } + } +#else + auto vinf = F16::set1(-INFINITY); + for (int l = 0; l < k_step/F16::block_size; ++l) { + vk[l] = apply_mask(l, mask, F16::load(cache + k_step*j + F16::block_size*l), vinf); + } + if (softcap <= 0) { + for (int l = 0; l < k_step/F16::block_size; ++l) vk[l] = F16::mul(vscale, vk[l]); + } else { + auto v_softcap = F16::set1(softcap); + for (int l = 0; l < k_step/F16::block_size; ++l) vk[l] = F16::mul(v_softcap, v_tanh(F16::mul(vscale, vk[l]))); + } +#endif + return F16::reduce_max(vk); + } +#endif + +#ifdef __aarch64__ + inline void update_M_S(int j, float32x4_t * vk) { + float smax = load_and_scale(j, vk); + update_M(j, smax); + if (M[j] > -INFINITY) update_S(j, vk); + } + inline void update_M_S(int j, float32x4_t * vk, const char * mask) { + float smax = load_apply_mask_and_scale(j, vk, mask); + update_M(j, smax); + if (M[j] > -INFINITY) update_S(j, vk); + } +#else + inline void update_M_S(int j, F16::Data * vk) { + float smax = load_and_scale(j, vk); + update_M(j, smax); + if (M[j] > -INFINITY) update_S(j, vk); + } + inline void update_M_S(int j, F16::Data * vk, const char * mask) { + float smax = load_apply_mask_and_scale(j, vk, mask); + update_M(j, smax); + if (M[j] > -INFINITY) update_S(j, vk); + } +#endif + + cache_t cache[q_step*k_step]; + float S[q_step], M[q_step]; + int need_scaling[q_step]; + float vms[q_step]; + const F16::Data vscale; + const float softcap; + const ggml_half h_inf; + +}; + +template +struct FlashQKV { + +#ifdef __aarch64__ + using qkv_cache_t = float16_t; +#else + using qkv_cache_t = float; +#endif + + template + inline void accumulate_qkv_1(const VHelper& vh, const FMS& fms) { + static_assert(q_step == FMS::q_step); + F16::Data vq[D/F16::block_size]; + if (fms.need_scaling[0] == 2) { + for (int i = 0; i < D/F16::block_size; ++i) vq[i] = F16::zero(); + } else { + for (int i = 0; i < D/F16::block_size; ++i) vq[i] = F16::load(qkv_cache + F16::block_size*i); + if (fms.need_scaling[0] == 1) { + auto vms = F16::set1(fms.vms[0]); + for (int i = 0; i < D/F16::block_size; ++i) vq[i] = F16::mul(vms, vq[i]); + } + } + F16::Data v0, v1; + for (int l = 0; l < k_step; l += 4) { + auto vs0 = F16::set1(fms.cache[l + 0]); + auto vs1 = F16::set1(fms.cache[l + 1]); + auto vs2 = F16::set1(fms.cache[l + 2]); + auto vs3 = F16::set1(fms.cache[l + 3]); + for (int i = 0; i < D/F16::block_size; i += 2) { + vh.load(l+0, i, v0, v1); + vq[i+0] = F16::fmadd(vq[i+0], v0, vs0); + vq[i+1] = F16::fmadd(vq[i+1], v1, vs0); + vh.load(l+1, i, v0, v1); + vq[i+0] = F16::fmadd(vq[i+0], v0, vs1); + vq[i+1] = F16::fmadd(vq[i+1], v1, vs1); + vh.load(l+2, i, v0, v1); + vq[i+0] = F16::fmadd(vq[i+0], v0, vs2); + vq[i+1] = F16::fmadd(vq[i+1], v1, vs2); + vh.load(l+3, i, v0, v1); + vq[i+0] = F16::fmadd(vq[i+0], v0, vs3); + vq[i+1] = F16::fmadd(vq[i+1], v1, vs3); + } + } + for (int i = 0; i < D/F16::block_size; ++i) F16::store(qkv_cache + F16::block_size*i, vq[i]); + } + + // This fails for head sizes of 80 and 112 as D/16 is odd, so we cannot do steps of 2 + // Hence, for now, we will not handle head sizes of 80 and 112 + template + inline void accumulate_qkv(const VHelper& vh, const FMS& fms) { + static_assert(q_step == FMS::q_step); + if constexpr (q_step == 1) { + accumulate_qkv_1(vh, fms); + return; + } + for (int j = 0; j < q_step; ++j) { + auto R = qkv_cache + D*j; + if (fms.need_scaling[j] == 2) { + std::memset(R, 0, D*sizeof(qkv_cache_t)); + } + else if (fms.need_scaling[j] == 1) { + auto vms = F16::set1(fms.vms[j]); + for (int i = 0; i < D/F16::block_size; ++i) { + F16::store(R + F16::block_size*i, F16::mul(vms, F16::load(R + F16::block_size*i))); + } + } + } +#ifdef __AVX512F__ + if constexpr ((D/F16::block_size)%4 == 0) { + F16::Data v[16]; + F16::Data vs[4]; + for (int i = 0; i < D/F16::block_size; i += 4) { + for (int l = 0; l < k_step; l += 4) { + for (int k = 0; k < 4; ++k) { + vh.load(l+k, i+0, v[4*k+0], v[4*k+1]); + vh.load(l+k, i+2, v[4*k+2], v[4*k+3]); + } + for (int j = 0; j < q_step; ++j) { + auto R = qkv_cache + D*j; + auto s1 = F16::load(R + F16::block_size*(i+0)); + auto s2 = F16::load(R + F16::block_size*(i+1)); + auto s3 = F16::load(R + F16::block_size*(i+2)); + auto s4 = F16::load(R + F16::block_size*(i+3)); + F16::set4(fms.cache + k_step*j + l, vs); + for (int k = 0; k < 4; ++k) { + s1 = F16::fmadd(s1, v[4*k+0], vs[k]); + s2 = F16::fmadd(s2, v[4*k+1], vs[k]); + s3 = F16::fmadd(s3, v[4*k+2], vs[k]); + s4 = F16::fmadd(s4, v[4*k+3], vs[k]); + } + F16::store(R + F16::block_size*(i+0), s1); + F16::store(R + F16::block_size*(i+1), s2); + F16::store(R + F16::block_size*(i+2), s3); + F16::store(R + F16::block_size*(i+3), s4); + } + } + } + return; + } +#endif + F16::Data v[8]; +#ifdef __AVX2__ + F16::Data vs[4]; +#endif + for (int i = 0; i < D/F16::block_size; i += 2) { + for (int l = 0; l < k_step; l += 4) { + vh.load(l+0, i, v[0], v[4]); + vh.load(l+1, i, v[1], v[5]); + vh.load(l+2, i, v[2], v[6]); + vh.load(l+3, i, v[3], v[7]); + for (int j = 0; j < q_step; ++j) { + auto R = qkv_cache + D*j; + auto s1 = F16::load(R + F16::block_size*(i+0)); + auto s2 = F16::load(R + F16::block_size*(i+1)); +#ifdef __AVX2__ + F16::set4(fms.cache + k_step*j + l, vs); + for (int k = 0; k < 4; ++k) { + s1 = F16::fmadd(s1, v[k+0], vs[k]); + s2 = F16::fmadd(s2, v[k+4], vs[k]); + } +#else + auto vs = F16::set4(fms.cache + k_step*j + l); + s1 = F16::fmadd_lane0(s1, v[0], vs); + s2 = F16::fmadd_lane0(s2, v[4], vs); + s1 = F16::fmadd_lane1(s1, v[1], vs); + s2 = F16::fmadd_lane1(s2, v[5], vs); + s1 = F16::fmadd_lane2(s1, v[2], vs); + s2 = F16::fmadd_lane2(s2, v[6], vs); + s1 = F16::fmadd_lane3(s1, v[3], vs); + s2 = F16::fmadd_lane3(s2, v[7], vs); +#endif + F16::store(R + F16::block_size*(i+0), s1); + F16::store(R + F16::block_size*(i+1), s2); + } + } + } + } + + template + inline void accumulate_qkv(int nq1, const VHelper& vh, const FMS& fms) { + static_assert(q_step == FMS::q_step); + if (nq1 == 1) { + accumulate_qkv_1(vh, fms); + return; + } + F16::Data v[8]; + for (int j = 0; j < nq1; ++j) { + auto R = qkv_cache + D*j; + if (fms.need_scaling[j] == 2) { + std::memset(R, 0, D*sizeof(qkv_cache_t)); + } + else if (fms.need_scaling[j] == 1) { + auto vms = F16::set1(fms.vms[j]); + for (int i = 0; i < D/F16::block_size; ++i) { + F16::store(R + F16::block_size*i, F16::mul(vms, F16::load(R + F16::block_size*i))); + } + } + } + for (int i = 0; i < D/F16::block_size; i += 2) { + for (int l = 0; l < k_step; l += 4) { + vh.load(l+0, i, v[0], v[4]); + vh.load(l+1, i, v[1], v[5]); + vh.load(l+2, i, v[2], v[6]); + vh.load(l+3, i, v[3], v[7]); + for (int j = 0; j < nq1; ++j) { + auto R = qkv_cache + D*j; + auto s1 = F16::load(R + F16::block_size*(i+0)); + auto s2 = F16::load(R + F16::block_size*(i+1)); + auto vs = F16::set4(fms.cache + k_step*j + l); + s1 = F16::fmadd_lane0(s1, v[0], vs); + s2 = F16::fmadd_lane0(s2, v[4], vs); + s1 = F16::fmadd_lane1(s1, v[1], vs); + s2 = F16::fmadd_lane1(s2, v[5], vs); + s1 = F16::fmadd_lane2(s1, v[2], vs); + s2 = F16::fmadd_lane2(s2, v[6], vs); + s1 = F16::fmadd_lane3(s1, v[3], vs); + s2 = F16::fmadd_lane3(s2, v[7], vs); + F16::store(R + F16::block_size*(i+0), s1); + F16::store(R + F16::block_size*(i+1), s2); + } + } + } + } + + template + inline void normalize_and_store_1row(const FMS& fms, int j, qkv_cache_t * R, float * qkv, const float * sinkf) const { + static_assert(q_step == FMS::q_step); + float S = fms.S[j]; + if (sinkf) { + float s = *sinkf; + if (s > fms.M[j]) { + float m = expf(fms.M[j] - s); + auto vm = F16::set1(m); + for (int i = 0; i < D/F16::block_size; ++i) { + auto Ri = R + F16::block_size*i; + F16::store(Ri, F16::mul(vm, F16::load(Ri))); + } + S = S*m + 1; + } else { + S += expf(s - fms.M[j]); + } + } + GGML_ASSERT(S > 0); + auto norm = F16::set1(1/S); + //auto norm = F16::set1(fms.S[j] > 0 ? 1/fms.S[j] : 0.f); + for (int i = 0; i < D/F16::block_size; ++i) { + auto r = F16::load(R + F16::block_size*i); + F16::store(qkv + F16::block_size*i, F16::mul(norm, r)); + } + } + + template + inline void normalize_and_store(const FMS& fms, int nq1, int stride_qkv, float * qkv, const float * sinkf, float * M, float * S) { + static_assert(q_step == FMS::q_step); + if (M && S) { + std::memcpy(M, fms.M, nq1*sizeof(float)); + std::memcpy(S, fms.S, nq1*sizeof(float)); + auto R = qkv_cache; + for (int j = 0; j < nq1; ++j) { +#ifdef __aarch64__ + for (int i = 0; i < D/F16::block_size; ++i) { + F16::store(qkv + F16::block_size*i, F16::load(R + F16::block_size*i)); + } +#else + std::memcpy(qkv, R, D*sizeof(float)); +#endif + qkv += stride_qkv; + R += D; + } + } else { + auto R = qkv_cache; + for (int j = 0; j < nq1; ++j) { + normalize_and_store_1row(fms, j, R, qkv, sinkf); + qkv += stride_qkv; + R += D; + } + } + } + + template + inline void normalize_and_store(const FMS& fms, int stride_qkv, float * qkv, const float * sinkf, float * M, float * S) { + static_assert(q_step == FMS::q_step); + if (M && S) { + std::memcpy(M, fms.M, q_step*sizeof(float)); + std::memcpy(S, fms.S, q_step*sizeof(float)); + auto R = qkv_cache; + for (int j = 0; j < q_step; ++j) { +#ifdef __aarch64__ + for (int i = 0; i < D/F16::block_size; ++i) { + F16::store(qkv + F16::block_size*i, F16::load(R + F16::block_size*i)); + } +#else + std::memcpy(qkv, R, D*sizeof(float)); +#endif + qkv += stride_qkv; + R += D; + } + } else { + auto R = qkv_cache; + for (int j = 0; j < q_step; ++j) { + normalize_and_store_1row(fms, j, R, qkv, sinkf); + qkv += stride_qkv; + R += D; + } + } + } + + // qkv_cache_t qkv_cache[D*q_step]; + // The initializer is not actually required. But the compiler cannot figure out that when qkv_cache is + // first used for q_step rows, fms.need_scaling[j] is always 2, which zeroes the content of qkv_cache. + // As a result, we get an infinite stream of warnings about uninitialized variable use (one for each + // combination of D, q_step, k_step), which is extremely annoying. Hence, I succumb to the trend of + // constantly being saved by others (the compiler in this case), and add this 100% unnecessary initialization. + qkv_cache_t qkv_cache[D*q_step]; // = {}; + //qkv_cache_t * qkv_cache; +}; + +template +struct FlashQKfp32 { + static_assert(D%F16::block_size == 0 && D <= 576); + static_assert(k_step%F16::block_size == 0); + static_assert(q_step <= 4 || q_step%4 == 0); + + template + static inline void multiply_mask_kq(const KHelper& kh, int stride_q, int stride_m, const q_float * q, const char * mask, + FlashMS& fms) { +#ifdef __AVX2__ + constexpr int nrc_k = 8; + static_assert(k_step%nrc_k == 0); +#endif + DataInfo info{fms.cache, (const char *)q, k_step, stride_q*sizeof(q_float), 0, 1, nullptr}; + iqk_gemm_default_floats(D, q_step, kh.block, kh.stride, info, k_step); +#ifdef __AVX2__ + F16::Data vk[k_step/F16::block_size]; +#else + float32x4_t vk[k_step/4]; +#endif + for (int j = 0; j < q_step; ++j) { + fms.update_M_S(j, vk, mask + stride_m*j); + } + } + + template + static inline void multiply_mask_kq(int nq, const KHelper& kh, int stride_q, int stride_m, const q_float * q, const char * mask, + FlashMS& fms) { +#ifdef __AVX2__ + constexpr int nrc_k = 8; + static_assert(k_step%nrc_k == 0); +#endif + DataInfo info{fms.cache, (const char *)q, k_step, stride_q*sizeof(q_float), 0, 1, nullptr}; + iqk_gemm_default_floats(D, nq, kh.block, kh.stride, info, k_step); +#ifdef __AVX2__ + F16::Data vk[k_step/F16::block_size]; +#else + float32x4_t vk[k_step/4]; +#endif + for (int j = 0; j < nq; ++j) { + fms.update_M_S(j, vk, mask + stride_m*j); + } + } + +#ifdef __aarch64__ + static inline void convert(int nq, int stride_q, const float * q, float16_t * q_f16) { + for (int i = 0; i < nq; ++i) { + for (int j = 0; j < D; j += 8) { + auto val1_f32 = vld1q_f32(q + j + 0); + auto val2_f32 = vld1q_f32(q + j + 4); + auto val_f16 = vcombine_f16(vcvt_f16_f32(val1_f32), vcvt_f16_f32(val2_f32)); + vst1q_f16(q_f16 + j, val_f16); + } + q += stride_q; + q_f16 += D; + } + } +#endif + + template + static inline void mul_mask_kq(const KHelper& kh, int stride_m, + const block_q8 * q, const char * mask, FlashMS& fms) { + // As far as I can tell, this static assert is a remnant of the times where the matrix multiplications were done inline + // here with bespoke kernels instead of just using the regular mat mul kernels. But, just in case, leaving it in place + // but commneted out. + //constexpr int kMaxQ = 8; + //static_assert(q_step < kMaxQ || q_step%kMaxQ == 0); + DataInfo info{fms.cache, (const char *)q, k_step, (D/KHelper::block_size_q)*sizeof(block_q8), 0, 1, nullptr}; + if constexpr (std::is_same_v> || + std::is_same_v>) { + iqk_gemm_q8kv_fa(D, q_step, kh.type, kh.block, kh.stride, info, k_step); + } else { + iqk_gemm_legacy_fa(D, q_step, kh.type, kh.block, kh.stride, info, k_step); + } +#ifdef __aarch64__ + float32x4_t vk[k_step/4]; + for (int j = 0; j < q_step; ++j) { + fms.update_M_S(j, vk, mask + stride_m*j); + } +#else + F16::Data vk[k_step/F16::block_size]; + for (int j = 0; j < q_step; ++j) { + fms.update_M_S(j, vk, mask + stride_m*j); + } +#endif + } + + template + static inline void mul_mask_kq(int nq, const KHelper& kh, int stride_m, + const block_q8 * q, const char * mask, FlashMS& fms) { + GGML_ASSERT(nq < q_step); + DataInfo info{fms.cache, (const char *)q, k_step, (D/KHelper::block_size_q)*sizeof(block_q8), 0, 1, nullptr}; + if constexpr (std::is_same_v> || + std::is_same_v>) { + iqk_gemm_q8kv_fa(D, nq, kh.type, kh.block, kh.stride, info, k_step); + } else { + iqk_gemm_legacy_fa(D, nq, kh.type, kh.block, kh.stride, info, k_step); + } +#ifdef __aarch64__ + float32x4_t vk[k_step/4]; + for (int j = 0; j < nq; ++j) { + fms.update_M_S(j, vk, mask + stride_m*j); + } +#else + F16::Data vk[k_step/F16::block_size]; + for (int j = 0; j < nq; ++j) { + fms.update_M_S(j, vk, mask + stride_m*j); + } +#endif + } +}; + +template +void compute_helper(KHelper& kh, VHelper& vh, int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, + FlashMS& fms, + FlashQKV& fqkv, + const float * q, const char * mask, float * qkv, + const float * sinkf, float * M, float * S) { +#ifdef __aarch64__ + float16_t q_f16[Dk*q_step]; +#endif + + for (int i1 = 0; i1 < nq1/q_step; ++i1) { + fms.init_qstep(); + kh.reset_block(); + vh.reset_block(); +#ifdef __aarch64__ + KQHelper::convert(q_step, stride_q, q, q_f16); +#endif + auto mr = mask; + int nk1_eff = mask_effective_nk1(mr, q_step, stride_m, nk1, k_step); + for (int k1 = 0; k1 < nk1_eff/k_step; ++k1) { +#ifdef __aarch64__ + KQHelper::multiply_mask_kq(kh, Dk, stride_m, q_f16, mr, fms); +#else + KQHelper::multiply_mask_kq(kh, stride_q, stride_m, q, mr, fms); +#endif + fqkv.accumulate_qkv(vh, fms); + kh.next_block(k_step); + vh.next_block(k_step); + mr += k_step*sizeof(ggml_half); + } + fqkv.normalize_and_store(fms, stride_qkv, qkv, sinkf, M, S); + + q += q_step*stride_q; + mask += q_step*stride_m; + qkv += q_step*stride_qkv; + if (M && S) { M += q_step; S += q_step; } + } + int n_left = nq1 - q_step*(nq1/q_step); + if (n_left > 0) { + fms.init_qstep(); + kh.reset_block(); + vh.reset_block(); +#ifdef __aarch64__ + KQHelper::convert(n_left, stride_q, q, q_f16); +#endif + auto mr = mask; + for (int k1 = 0; k1 < nk1/k_step; ++k1) { +#ifdef __aarch64__ + KQHelper::multiply_mask_kq(n_left, kh, Dk, stride_m, q_f16, mr, fms); +#else + KQHelper::multiply_mask_kq(n_left, kh, stride_q, stride_m, q, mr, fms); +#endif + fqkv.accumulate_qkv(n_left, vh, fms); + kh.next_block(k_step); + vh.next_block(k_step); + mr += k_step*sizeof(ggml_half); + } + fqkv.normalize_and_store(fms, n_left, stride_qkv, qkv, sinkf, M, S); + } +} + +template +void compute_helper_q(KHelper& kh, VHelper& vh, int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, + FlashMS& fms, + FlashQKV& fqkv, + const float * q, const char * mask, float * qkv, + const float * sinkf, float * M, float * S, char * qptr) { + auto q8 = (typename KHelper::block_q8 *)qptr; + // This optimization fails under certain conditions (see https://github.com/ikawrakow/ik_llama.cpp/issues/1205) + // => disabling until I figure out what goes wrong + if constexpr (q_step >= 4 && std::is_same_v) { + if (nq1 == q_step) { + fms.init_qstep(); + kh.reset_block(); + vh.reset_block(); + block_q8_0_r8 q8r8[Dk/QK8_0 * k_step/8]; + HelperQ80R8 khr8((const char *)q8r8, Dk/QK8_0*sizeof(block_q8_0)); + auto q8r = (typename HelperQ80R8::block_q8 *)qptr; + HelperQ80::convert(q_step, stride_q, q, q8r); + auto mr = mask; + int nk1_eff = mask_effective_nk1(mr, q_step, stride_m, nk1, k_step); + for (int k1 = 0; k1 < nk1_eff/k_step; ++k1) { + HelperQ80R8::repack(k_step, kh.block, kh.stride, q8r8); + KQHelper::mul_mask_kq(khr8, stride_m, q8r, mr, fms); + fqkv.accumulate_qkv(vh, fms); + kh.next_block(k_step); + vh.next_block(k_step); + mr += k_step*sizeof(ggml_half); + } + fqkv.normalize_and_store(fms, stride_qkv, qkv, sinkf, M, S); + return; + } + } +#if FA_TIMING + Perf perf(false); +#endif + for (int i1 = 0; i1 < nq1/q_step; ++i1) { +#if FA_TIMING + auto t1 = Perf::cur_time(); +#endif + fms.init_qstep(); + kh.reset_block(); + vh.reset_block(); + HelperQ80::convert(q_step, stride_q, q, q8); +#if FA_TIMING + perf.accum_nolock(0, t1); +#endif + auto mr = mask; + int nk1_eff = mask_effective_nk1(mr, q_step, stride_m, nk1, k_step); + for (int k1 = 0; k1 < nk1_eff/k_step; ++k1) { +#if FA_TIMING + t1 = Perf::cur_time(); + KQHelper::mul_mask_kq(kh, stride_m, q8, mr, fms); + perf.accum_nolock(1, t1); + t1 = Perf::cur_time(); + fqkv.accumulate_qkv(vh, fms); + perf.accum_nolock(2, t1); +#else + KQHelper::mul_mask_kq(kh, stride_m, q8, mr, fms); + fqkv.accumulate_qkv(vh, fms); +#endif + kh.next_block(k_step); + vh.next_block(k_step); + mr += k_step*sizeof(ggml_half); + } +#if FA_TIMING + t1 = Perf::cur_time(); + fqkv.normalize_and_store(fms, stride_qkv, qkv, sinkf, M, S); + perf.accum_nolock(3, t1); +#else + fqkv.normalize_and_store(fms, stride_qkv, qkv, sinkf, M, S); +#endif + + q += q_step*stride_q; + mask += q_step*stride_m; + qkv += q_step*stride_qkv; + if (M && S) { M += q_step; S += q_step; } + } + int n_left = nq1 - q_step*(nq1/q_step); + if (n_left > 0) { + fms.init_qstep(); + kh.reset_block(); + vh.reset_block(); + HelperQ80::convert(n_left, stride_q, q, q8); + auto mr = mask; + for (int k1 = 0; k1 < nk1/k_step; ++k1) { + KQHelper::mul_mask_kq(n_left, kh, stride_m, q8, mr, fms); + fqkv.accumulate_qkv(n_left, vh, fms); + kh.next_block(k_step); + vh.next_block(k_step); + mr += k_step*sizeof(ggml_half); + } + fqkv.normalize_and_store(fms, n_left, stride_qkv, qkv, sinkf, M, S); + } +#if FA_TIMING + Perf::instance().add(perf); +#endif +} + +char * get_q_storage(size_t size) { + thread_local std::vector q_storage; + if (q_storage.size() < size) q_storage.resize(size); + return q_storage.data(); +} + +// Some of the methods in FlashAttn have two identical implementations that only differ by +// one version using a loop over the template parameter q_step, while the other using a loop +// over an input parameter nq (these are loops over the rows of q^T). I dislike this a lot, +// but performance drops signficantly if I remove the version with fixed q_step iterations. +// We only instantiate FlashAttn with q_step = 1 and q_step = 4 or 8 (depending on head size D), +// so when we have to process Nq rows, we process q_step*(Nq/q_step) using fixed q_step loops, +// and use the variable nq version (with lower performance) only for the remaining i1...q_step-1 +// rows (if Nq is not a multiple of q_step). One could have made the number of q^T rows to +// process template parameter of such functions, but this would result in the compiler generating +// q_step-1 versions of these functions for us, which I though was too much with q_step = 8. +template +struct FlashAttn { + static_assert(Dk%F16::block_size == 0 && Dk <= 576); + static_assert(Dv%F16::block_size == 0 && Dv <= 512); + static_assert(k_step%F16::block_size == 0); + static_assert(q_step <= 4 || q_step%4 == 0); + + FlashAttn(float scale, float softcap, const float * sinkf) : fms(scale, softcap), sinkf(sinkf) {} + + template + void compute(KHelper& kh, VHelper& vh, int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, + const float * q, const char * mask, float * qkv, [[maybe_unused]] float * M, [[maybe_unused]] float * S) { + if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v> || + std::is_same_v || + std::is_same_v> || + std::is_same_v>) { + constexpr size_t kMaxOnStackSize = 576; + //auto q_size = q_step*(Dk/KHelper::block_size_q)*sizeof(typename KHelper::block_q8); + auto q_size = q_step*(Dk/QK8_2*sizeof(block_q8_2)); + q_size = GGML_PAD(q_size, 64); + if (q_size > kMaxOnStackSize) { + auto qptr = get_q_storage(q_size); + if (false && nq1 >= 8) { + if constexpr (std::is_same_v) { +#if FA_TIMING + auto t1 = Perf::cur_time(); + HelperQ80R8 khr4(nk1, kh); + Perf::instance().accum(4, t1); +#else + HelperQ80R8 khr4(nk1, kh); +#endif + compute_helper_q, VHelper, FlashQKfp32>( + khr4, vh, nq1, nk1, stride_q, stride_m, stride_qkv, fms, fqkv, q, mask, qkv, sinkf, M, S, qptr); + return; + + } +#if GGML_IQK_FA_ALL_QUANTS + if constexpr (std::is_same_v>) { +#if FA_TIMING + auto t1 = Perf::cur_time(); + HelperQ8KVR8 khr4(nk1, kh); + Perf::instance().accum(4, t1); +#else + HelperQ8KVR8 khr4(nk1, kh); +#endif + compute_helper_q, VHelper, FlashQKfp32>( + khr4, vh, nq1, nk1, stride_q, stride_m, stride_qkv, fms, fqkv, q, mask, qkv, sinkf, M, S, qptr); + return; + } +#endif + } + compute_helper_q>( + kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, fms, fqkv, q, mask, qkv, sinkf, M, S, qptr); + + } + else { + typename KHelper::block_q8 q8[q_step*(Dk/KHelper::block_size_q)]; + compute_helper_q>( + kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, fms, fqkv, q, mask, qkv, sinkf, M, S, (char *)q8); + } + } + else { + compute_helper>( + kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, fms, fqkv, q, mask, qkv, sinkf, M, S); + } + } + + FlashMS fms; + FlashQKV fqkv; + const float * sinkf; + +}; + +#ifdef __AVX512BF16__ + +template +struct HelperBF16 final : public BaseHelper { + using Base = BaseHelper; + HelperBF16(const char * data, int stride) : Base(data, stride) {} + inline void load(int l1, __m512bh * vk) const { + auto dr = Base::lblock(l1); + for (int i = 0; i < D/32; ++i) vk[i] = __m512bh(_mm512_loadu_si512((const __m512i*)dr + i)); + } + + inline void load(int l1, int i, __m512& v1, __m512& v2) const { + auto dr = Base::lblock(l1); + v1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)dr + i + 0)), 16)); + v2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i *)dr + i + 1)), 16)); + } + + inline void load_2(int l1, __m512bh * vk) const { + load(l1+0, vk+0); + load(l1+1, vk+D/32); + } + + inline void load_4(int l1, __m512bh * vk) const { + load(l1+0, vk+0); + load(l1+1, vk+1*D/32); + load(l1+2, vk+2*D/32); + load(l1+3, vk+3*D/32); + } + + inline void load_8(int l1, __m512bh * vk) const { + for (int k = 0; k < 8; ++k) load(l1 + k, vk + k*D/32); + } +}; + +template +struct FlashQKbf16 { + //static_assert(D%32 == 0 && D <= 256); + static_assert(D%32 == 0 && D <= 576); + static_assert(k_step%32 == 0); + static_assert(q_step <= 4 || q_step%4 == 0); + + static inline void mult_mask_kq_one(int l1, int m1, int stride_q, int stride_m, const float * q, const char * mask, + __m512bh * qv, const __m512bh * vkh, FlashMS& fms) { + // q index is q_step*i1 + m1 + // k index is k_step*k1 + l1 + const ggml_half * mp = (const ggml_half *)(mask + stride_m*m1); + fms.cache[k_step*m1 + l1 + 0] = fms.cache[k_step*m1 + l1 + 1] = -INFINITY; + if (mp[l1+0] == fms.h_inf && mp[l1+1] == fms.h_inf) { + return; + } + auto qr = q + m1*stride_q; + for (int i = 0; i < D/32; ++i) { + auto val1 = _mm512_loadu_ps(qr + 32*i); + auto val2 = _mm512_loadu_ps(qr + 32*i + 16); + qv[i] = _mm512_cvtne2ps_pbh(val2, val1); + } + if (mp[l1+0] != fms.h_inf) { + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i], qv[i]); + fms.cache[k_step*m1 + l1 + 0] = _mm512_reduce_add_ps(vsum); + } + if (mp[l1+1] != fms.h_inf) { + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i+D/32], qv[i]); + fms.cache[k_step*m1 + l1 + 1] = _mm512_reduce_add_ps(vsum); + } + } + + static inline void mult_mask_kq_one(int l1, int m1, int stride_m, const ggml_bf16_t * q, const char * mask, + __m512bh * qv, const __m512bh * vkh, FlashMS& fms) { + // q index is q_step*i1 + m1 + // k index is k_step*k1 + l1 + const ggml_half * mp = (const ggml_half *)(mask + stride_m*m1); + fms.cache[k_step*m1 + l1 + 0] = fms.cache[k_step*m1 + l1 + 1] = -INFINITY; + if (mp[l1+0] == fms.h_inf && mp[l1+1] == fms.h_inf) { + return; + } + auto qr = q + m1*D; + for (int i = 0; i < D/32; ++i) qv[i] = __m512bh(_mm512_loadu_si512((const __m512i*)qr + i)); + if (mp[l1+0] != fms.h_inf) { + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i], qv[i]); + fms.cache[k_step*m1 + l1 + 0] = _mm512_reduce_add_ps(vsum); + } + if (mp[l1+1] != fms.h_inf) { + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i+D/32], qv[i]); + fms.cache[k_step*m1 + l1 + 1] = _mm512_reduce_add_ps(vsum); + } + } + + static inline void mult_mask_kq_4(int l1, int m1, int stride_q, int stride_m, const float * q, const char * mask, + __m512bh * qv, const __m512bh * vkh, FlashMS& fms) { + // q index is q_step*i1 + m1 + // k index is k_step*k1 + l1 + const ggml_half * mp = (const ggml_half *)(mask + stride_m*m1); + fms.cache[k_step*m1 + l1 + 0] = fms.cache[k_step*m1 + l1 + 1] = + fms.cache[k_step*m1 + l1 + 2] = fms.cache[k_step*m1 + l1 + 3] = -INFINITY; + if (mp[l1+0] == fms.h_inf && mp[l1+1] == fms.h_inf && mp[l1+2] == fms.h_inf && mp[l1+3] == fms.h_inf) { + return; + } + auto qr = q + m1*stride_q; + for (int i = 0; i < D/32; ++i) { + auto val1 = _mm512_loadu_ps(qr + 32*i); + auto val2 = _mm512_loadu_ps(qr + 32*i + 16); + qv[i] = _mm512_cvtne2ps_pbh(val2, val1); + } + for (int k = 0; k < 4; ++k) { + if (mp[l1+k] == fms.h_inf) continue; + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i+k*(D/32)], qv[i]); + fms.cache[k_step*m1 + l1 + k] = _mm512_reduce_add_ps(vsum); + } + } + + static inline void mult_mask_kq_4(int l1, int m1, int stride_m, const ggml_bf16_t * q, const char * mask, + __m512bh * qv, const __m512bh * vkh, FlashMS& fms) { + // q index is q_step*i1 + m1 + // k index is k_step*k1 + l1 + const ggml_half * mp = (const ggml_half *)(mask + stride_m*m1); + fms.cache[k_step*m1 + l1 + 0] = fms.cache[k_step*m1 + l1 + 1] = + fms.cache[k_step*m1 + l1 + 2] = fms.cache[k_step*m1 + l1 + 3] = -INFINITY; + if (mp[l1+0] == fms.h_inf && mp[l1+1] == fms.h_inf && mp[l1+2] == fms.h_inf && mp[l1+3] == fms.h_inf) { + return; + } + auto qr = q + m1*D; + for (int i = 0; i < D/32; ++i) qv[i] = __m512bh(_mm512_loadu_si512((const __m512i *)qr + i)); + for (int k = 0; k < 4; ++k) { + if (mp[l1+k] == fms.h_inf) continue; + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i+k*(D/32)], qv[i]); + fms.cache[k_step*m1 + l1 + k] = _mm512_reduce_add_ps(vsum); + } + } + + static inline __m128 hsum_float_4x4(__m128 * a) { + for (int i = 0; i < 2; ++i) a[i] = _mm_add_ps(_mm_unpacklo_ps(a[i], a[i+2]), _mm_unpackhi_ps(a[i], a[i+2])); + return _mm_add_ps(_mm_unpacklo_ps(a[0], a[1]), _mm_unpackhi_ps(a[0], a[1])); + } + + template + static inline void multiply_mask_kq(const KHelper& kh, int stride_q, int stride_m, const float * q, + const char * mask, FlashMS& fms) { + { + __m512bh qv[D/32]; + if constexpr (D <= 128) { + __m512bh vkh[D/8]; + for (int l1 = 0; l1 < k_step; l1 += 4) { + kh.load_4(l1, vkh); + for (int j = 0; j < q_step; ++j) { + mult_mask_kq_4(l1, j, stride_q, stride_m, q, mask, qv, vkh, fms); + } + } + } else { + __m512bh vkh[D/16]; + for (int l1 = 0; l1 < k_step; l1 += 2) { + kh.load_2(l1, vkh); + for (int j = 0; j < q_step; ++j) { + mult_mask_kq_one(l1, j, stride_q, stride_m, q, mask, qv, vkh, fms); + } + } + } + } + __m512 vk[k_step/16]; + for (int j = 0; j < q_step; ++j) { + fms.update_M_S(j, vk); + } + } + + static inline void mult_mask_kq_4(int l1, int m1, const ggml_bf16_t * q, + __m512bh * qv, const __m512bh * vkh, FlashMS& fms) { + auto qr = q + m1*D; + for (int i = 0; i < D/32; ++i) qv[i] = __m512bh(_mm512_loadu_si512((const __m512i *)qr + i)); + __m128 sum[4]; + for (int k = 0; k < 4; ++k) { + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i+k*(D/32)], qv[i]); + auto aux = _mm256_add_ps(_mm512_castps512_ps256(vsum), _mm512_extractf32x8_ps(vsum, 1)); + sum[k] = _mm_add_ps(_mm256_castps256_ps128(aux), _mm256_extractf128_ps(aux, 1)); + } + //auto sum4 = _mm_mask_blend_ps(m8, hsum_float_4x4(sum), _mm_set1_ps(-INFINITY)); + //_mm_storeu_ps(fms.cache + k_step*m1 + l1, sum4); + _mm_storeu_ps(fms.cache + k_step*m1 + l1, hsum_float_4x4(sum)); + } + + static IQK_ALWAYS_INLINE __m256 hsum_float_8x8(__m256 * accm) { + for (int i = 0; i < 4; ++i) { + accm[i] = _mm256_add_ps(_mm256_permute2f128_ps(accm[i], accm[i+4], 0x20), _mm256_permute2f128_ps(accm[i], accm[i+4], 0x31)); + //accm[i] = _mm256_set_m128(_mm_add_ps(_mm256_castps256_ps128(accm[i+4]), _mm256_extractf128_ps(accm[i+4], 1)), + // _mm_add_ps(_mm256_castps256_ps128(accm[i+0]), _mm256_extractf128_ps(accm[i+0], 1))); + } + for (int i = 0; i < 2; ++i) accm[i] = _mm256_add_ps(_mm256_unpacklo_ps(accm[i], accm[i+2]), _mm256_unpackhi_ps(accm[i], accm[i+2])); + return _mm256_add_ps(_mm256_unpacklo_ps(accm[0], accm[1]), _mm256_unpackhi_ps(accm[0], accm[1])); + } + + static inline void mult_mask_kq_8(int l1, int m1, const ggml_bf16_t * q, + __m512bh * qv, const __m512bh * vkh, FlashMS& fms) { + auto qr = q + m1*D; + for (int i = 0; i < D/32; ++i) qv[i] = __m512bh(_mm512_loadu_si512((const __m512i *)qr + i)); + __m256 sum[8]; + for (int k = 0; k < 8; ++k) { + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i+k*(D/32)], qv[i]); + sum[k] = _mm256_add_ps(_mm512_castps512_ps256(vsum), _mm512_extractf32x8_ps(vsum, 1)); + } + _mm256_storeu_ps(fms.cache + k_step*m1 + l1, hsum_float_8x8(sum)); + } + + static inline void mult_mask_kq_one(int l1, int m1, const ggml_bf16_t * q, + __m512bh * qv, const __m512bh * vkh, FlashMS& fms) { + auto qr = q + m1*D; + for (int i = 0; i < D/32; ++i) qv[i] = __m512bh(_mm512_loadu_si512((const __m512i*)qr + i)); + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i], qv[i]); + fms.cache[k_step*m1 + l1 + 0] = _mm512_reduce_add_ps(vsum); + vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vkh[i+D/32], qv[i]); + fms.cache[k_step*m1 + l1 + 1] = _mm512_reduce_add_ps(vsum); + } + +#if FA_TIMING + template + static inline void multiply_mask_kq(const KHelper& kh, int stride_m, const ggml_bf16_t * q, + const char * mask, FlashMS& fms, Perf& perf) { + auto t1 = Perf::cur_time(); +#else + template + static inline void multiply_mask_kq(const KHelper& kh, int stride_m, const ggml_bf16_t * q, + const char * mask, FlashMS& fms) { +#endif + if constexpr (q_step == 1) { + __m512bh vq[D/32]; + __m512bh vk[D/32]; + __m256 sum[8]; + for (int i = 0; i < D/32; ++i) vq[i] = __m512bh(_mm512_loadu_si512((const __m512i *)q + i)); + for (int l = 0; l < k_step; l += 8) { + for (int k = 0; k < 8; ++k) { + kh.load(l+k, vk); + auto vsum = _mm512_setzero_ps(); + for (int i = 0; i < D/32; ++i) vsum = _mm512_dpbf16_ps(vsum, vk[i], vq[i]); + sum[k] = _mm256_add_ps(_mm512_castps512_ps256(vsum), _mm512_extractf32x8_ps(vsum, 1)); + } + _mm256_storeu_ps(fms.cache + l, hsum_float_8x8(sum)); + } + } + else { + __m512bh qv[D/32]; + if constexpr (D <= 128) { + __m512bh vkh[D/4]; + for (int l1 = 0; l1 < k_step; l1 += 8) { + kh.load_8(l1, vkh); + for (int j = 0; j < q_step; ++j) mult_mask_kq_8(l1, j, q, qv, vkh, fms); + } + } else { + __m512bh vkh[D/16]; + for (int l1 = 0; l1 < k_step; l1 += 2) { + kh.load_2(l1, vkh); + for (int j = 0; j < q_step; ++j) mult_mask_kq_one(l1, j, q, qv, vkh, fms); + } + } + } +#if FA_TIMING + perf.accum_nolock(1, t1); + t1 = Perf::cur_time(); +#endif + F16::Data vk[k_step/16]; + for (int j = 0; j < q_step; ++j) { + fms.update_M_S(j, vk, mask + stride_m*j); + } +#if FA_TIMING + perf.accum_nolock(2, t1); +#endif + } + + template + static inline void multiply_mask_kq(int nq, const KHelper& kh, int stride_m, const ggml_bf16_t * q, + const char * mask, FlashMS& fms) { + { + __m512bh qv[D/32]; + if constexpr (D <= 128) { + __m512bh vkh[D/8]; + for (int l1 = 0; l1 < k_step; l1 += 4) { + kh.load_4(l1, vkh); + for (int j = 0; j < nq; ++j) mult_mask_kq_4(l1, j, q, qv, vkh, fms); + } + } else { + __m512bh vkh[D/16]; + for (int l1 = 0; l1 < k_step; l1 += 2) { + kh.load_2(l1, vkh); + for (int j = 0; j < nq; ++j) mult_mask_kq_one(l1, j, q, qv, vkh, fms); + } + } + } + F16::Data vk[k_step/16]; + for (int j = 0; j < nq; ++j) { + fms.update_M_S(j, vk, mask + stride_m*j); + } + } + + template + static inline void multiply_mask_kq(int nq, const KHelper& kh, int stride_q, int stride_m, const float * q, + const char * mask, FlashMS& fms) { + { + __m512bh qv[D/32]; + __m512bh vkh[D/16]; + for (int l1 = 0; l1 < k_step; l1 += 2) { + kh.load_2(l1, vkh); + for (int m1 = 0; m1 < nq; ++m1) { + mult_mask_kq_one(l1, m1, stride_q, stride_m, q, mask, qv, vkh, fms); + } + } + } + __m512 vk[k_step/16]; + for (int j = 0; j < nq; ++j) { + fms.update_M_S(j, vk); + } + } + + static inline void convert(int stride_q, const float * q, ggml_bf16_t * bf16) { + auto qr = q; + for (int j = 0; j < q_step; ++j) { + for (int i = 0; i < D/32; ++i) { + auto val1 = _mm512_loadu_ps(qr + 32*i); + auto val2 = _mm512_loadu_ps(qr + 32*i + 16); + _mm512_storeu_si512((__m512i *)bf16 + i, (__m512i)_mm512_cvtne2ps_pbh(val2, val1)); + } + qr += stride_q; + bf16 += D; + } + } + + static inline void convert(int nq, int stride_q, const float * q, ggml_bf16_t * bf16) { + auto qr = q; + for (int j = 0; j < nq; ++j) { + for (int i = 0; i < D/32; ++i) { + auto val1 = _mm512_loadu_ps(qr + 32*i); + auto val2 = _mm512_loadu_ps(qr + 32*i + 16); + _mm512_storeu_si512((__m512i *)bf16 + i, (__m512i)_mm512_cvtne2ps_pbh(val2, val1)); + } + qr += stride_q; + bf16 += D; + } + } +}; + +template +struct FlashAttnBF16 { + //static_assert(Dk%32 == 0 && Dk <= 256); + //static_assert(Dv%32 == 0 && Dv <= 256); + static_assert(Dk%32 == 0 && Dk <= 576); + static_assert(Dv%32 == 0 && Dv <= 512); + static_assert(k_step%32 == 0); + static_assert(q_step <= 4 || q_step%4 == 0); + + FlashAttnBF16(float scale, float softcap, const float * sinkf) : fms(scale, softcap), sinkf(sinkf) {} + + template + void compute(KHelper& kh, VHelper& vh, int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, + const float * q, const char * mask, float * qkv, [[maybe_unused]] float * M, [[maybe_unused]] float * S) { + ggml_bf16_t q_bf16[q_step*Dk]; +#if FA_TIMING + Perf perf(false); +#endif + for (int i1 = 0; i1 < nq1/q_step; ++i1) { +#if FA_TIMING + auto t1 = Perf::cur_time(); +#endif + fms.init_qstep(); + kh.reset_block(); + vh.reset_block(); + FlashQKbf16::convert(stride_q, q, q_bf16); +#if FA_TIMING + perf.accum_nolock(0, t1); +#endif + auto mr = mask; + int nk1_eff = mask_effective_nk1(mr, q_step, stride_m, nk1, k_step); + for (int k1 = 0; k1 < nk1_eff/k_step; ++k1) { +#if FA_TIMING + //t1 = Perf::cur_time(); + FlashQKbf16::multiply_mask_kq(kh, stride_m, q_bf16, mr, fms, perf); + //perf.accum_nolock(1, t1); + t1 = Perf::cur_time(); + fqkv.accumulate_qkv(vh, fms); + perf.accum_nolock(3, t1); +#else + FlashQKbf16::multiply_mask_kq(kh, stride_m, q_bf16, mr, fms); + fqkv.accumulate_qkv(vh, fms); +#endif + kh.next_block(k_step); + vh.next_block(k_step); + mr += k_step*sizeof(ggml_half); + } +#if FA_TIMING + t1 = Perf::cur_time(); +#endif + fqkv.normalize_and_store(fms, stride_qkv, qkv, sinkf, M, S); +#if FA_TIMING + perf.accum_nolock(4, t1); +#endif + + q += q_step*stride_q; + mask += q_step*stride_m; + qkv += q_step*stride_qkv; + if (M && S) { M += q_step; S += q_step; } + } + int n_left = nq1 - q_step*(nq1/q_step); + if (n_left > 0) { + fms.init_qstep(); + kh.reset_block(); + vh.reset_block(); + FlashQKbf16::convert(n_left, stride_q, q, q_bf16); + auto mr = mask; + for (int k1 = 0; k1 < nk1/k_step; ++k1) { + FlashQKbf16::multiply_mask_kq(n_left, kh, stride_m, q_bf16, mr, fms); + fqkv.accumulate_qkv(n_left, vh, fms); + kh.next_block(k_step); + vh.next_block(k_step); + mr += k_step*sizeof(ggml_half); + } + fqkv.normalize_and_store(fms, n_left, stride_qkv, qkv, sinkf, M, S); + } +#if FA_TIMING + Perf::instance().add(perf); +#endif + } + + FlashMS fms; + FlashQKV fqkv; + const float * sinkf; +}; +#endif + +template +inline void iqk_flash_helper(KHelper& kh, VHelper& vh, int nq1, int nk1, int stride_q, int stride_m, int stride_qkv, + const float * q, const char * mask, float scale, float softcap, float * qkv, + const float * sinkf, float * M, float * S) { + + auto update = [&nq1, &mask, &q, &qkv, &M, &S, stride_q, stride_m, stride_qkv] (int n) { + nq1 -= n; + if (nq1 == 0) return true; + q += n*stride_q; + mask += n*stride_m; + qkv += n*stride_qkv; + if (M && S) { M += n; S += n; } + return false; + }; + if (nk1 >= 512) { + if (nq1 >= 128) { + int n_step = nq1/128; + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, 128*n_step, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + if (update(128*n_step)) return; + } + if (nq1 >= 64) { + int n_step = nq1/64; + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, 64*n_step, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + if (update(64*n_step)) return; + } + if (nq1 >= 32) { + int n_step = nq1/32; + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, 32*n_step, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + if (update(32*n_step)) return; + } + if (nq1 >= 16) { + int n_step = nq1/16; + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, 16*n_step, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + if (update(16*n_step)) return; + } + } + if (nq1 == 12) { + // Special case: TG for GLM-4.5/4.6 + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, 12, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + return; + } + if (nq1 >= 8) { + int n_step = nq1/8; + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, 8*n_step, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + if (update(8*n_step)) return; + } + else if (nq1 >= 4) { + int n_step = nq1/4; + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, 4*n_step, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + if (update(4*n_step)) return; + } + else if (nq1 >= 2) { + int n_step = nq1/2; + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, 2*n_step, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + if (update(2*n_step)) return; + } + FlashAttn fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); +} + +#ifdef __AVX512BF16__ +template +inline void iqk_flash_helper_T(int nq1, int nk1, int stride_q, int stride_k, int stride_v, int stride_m, int stride_qkv, + const float * q, const char * k, const char * v, const char * mask, + float scale, float softcap, float * qkv, const float * sinkf, float * M, float * S) { + HelperBF16 kh(k, stride_k); + HelperBF16 vh(v, stride_v); + if (nk1 >= 4096) { + if (nq1 >= 64) { + FlashAttnBF16 fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + return; + } + else if (nq1 >= 16) { + FlashAttnBF16 fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + return; + } + } + if (nq1 >= 8) { + FlashAttnBF16 fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + } else { + FlashAttnBF16 fa(scale, softcap, sinkf); + fa.compute(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, (const char *)mask, qkv, M, S); + } +} +#endif + +template +inline bool iqk_flash_helper_T(KHelper& kh, ggml_type type_v, + int nq1, int nk1, int stride_q, int stride_v, int stride_m, int stride_qkv, + const float * q, const char * v, const char * mask, + float scale, float softcap, float * qkv, const float * sinkf, float * M, float * S) { + + switch (type_v) { + case GGML_TYPE_F16: { + HelperF16 vh(v, stride_v); + iqk_flash_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + } break; +#ifdef __AVX512BF16__ + case GGML_TYPE_BF16: { + HelperBF16 vh(v, stride_v); + iqk_flash_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + } break; +#endif + case GGML_TYPE_Q8_0: { + HelperQ80 vh(v, stride_v); + iqk_flash_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_Q8_KV: { + HelperQ8KV vh(v, stride_v); + iqk_flash_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_Q6_0: { + HelperQ60 vh(v, stride_v); + iqk_flash_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + } break; +#if GGML_IQK_FA_ALL_QUANTS + case GGML_TYPE_Q4_0: { + HelperQ40 vh(v, stride_v); + iqk_flash_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_Q4_1: { + HelperQ41 vh(v, stride_v); + iqk_flash_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_IQ4_NL: { + HelperIQ4nl vh(v, stride_v); + iqk_flash_helper(kh, vh, nq1, nk1, stride_q, stride_m, stride_qkv, q, mask, scale, softcap, qkv, sinkf, M, S); + } break; +#endif + default: return false; + } + return true; +} + +template +inline bool iqk_flash_helper_T(ggml_type type_k, ggml_type type_v, + int nq1, int nk1, int stride_q, int stride_k, int stride_v, int stride_m, int stride_qkv, + const float * q, const char * k, const char * v, const char * mask, + float scale, float softcap, float * qkv, const float * sinkf, float * M, float * S) { + + bool result = false; + switch (type_k) { + case GGML_TYPE_F16: { + HelperF16 kh(k, stride_k); + result = iqk_flash_helper_T(kh, type_v, nq1, nk1, stride_q, stride_v, stride_m, stride_qkv, q, v, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_Q8_0: { + HelperQ80 kh(k, stride_k); + result = iqk_flash_helper_T(kh, type_v, nq1, nk1, stride_q, stride_v, stride_m, stride_qkv, q, v, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_Q8_0_R8: { + HelperQ80R8 kh(k, stride_k); + result = iqk_flash_helper_T(kh, type_v, nq1, nk1, stride_q, stride_v, stride_m, stride_qkv, q, v, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_Q6_0: { + HelperQ60 kh(k, stride_k); + result = iqk_flash_helper_T(kh, type_v, nq1, nk1, stride_q, stride_v, stride_m, stride_qkv, q, v, mask, scale, softcap, qkv, sinkf, M, S); + } break; +#if GGML_IQK_FA_ALL_QUANTS + case GGML_TYPE_Q8_KV: { + HelperQ8KV kh(k, stride_k); + result = iqk_flash_helper_T(kh, type_v, nq1, nk1, stride_q, stride_v, stride_m, stride_qkv, q, v, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_Q4_0: { + HelperQ40 kh(k, stride_k); + result = iqk_flash_helper_T(kh, type_v, nq1, nk1, stride_q, stride_v, stride_m, stride_qkv, q, v, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_Q4_1: { + HelperQ41 kh(k, stride_k); + result = iqk_flash_helper_T(kh, type_v, nq1, nk1, stride_q, stride_v, stride_m, stride_qkv, q, v, mask, scale, softcap, qkv, sinkf, M, S); + } break; + case GGML_TYPE_IQ4_NL: { + HelperIQ4nl kh(k, stride_k); + result = iqk_flash_helper_T(kh, type_v, nq1, nk1, stride_q, stride_v, stride_m, stride_qkv, q, v, mask, scale, softcap, qkv, sinkf, M, S); + } break; +#endif + default: break; + } + + return result; +} + +} + +#define IQK_FA_CASE(name) bool name(int int_type_k, int int_type_v,int nq,int nk,\ + int stride_q, int stride_k, int stride_v, int stride_m, int stride_qkv,\ + const float * q, const void * k, const void * v, const void * mask,\ + float scale, float softcap,\ + float * qkv, const float * sinkf, float * M, float * S) + +IQK_FA_CASE(iqk_fa_576_512); +IQK_FA_CASE(iqk_fa_512_512); +IQK_FA_CASE(iqk_fa_320_256); +IQK_FA_CASE(iqk_fa_192_128); +IQK_FA_CASE(iqk_fa_192_192); +IQK_FA_CASE(iqk_fa_256_256); +IQK_FA_CASE(iqk_fa_128_128); +IQK_FA_CASE(iqk_fa_96_96); +IQK_FA_CASE(iqk_fa_64_64); + +#endif + diff --git a/llama.cpp/ggml/src/iqk/ggml-common.h b/llama.cpp/ggml/src/iqk/ggml-common.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/ggml-common.h @@ -0,0 +1,13 @@ +// opencoti F5-opt W2 (#290) — iqk-private ggml-common.h shim (HARD ISOLATION). +// +// iqk TUs resolve `#include "ggml-common.h"` to THIS file (their own dir wins +// over -Iggml/src). It forwards to llamafile's real, re-includable ggml-common.h +// (so shared/GGUF-frozen types stay SINGLE-SOURCED — they reach every TU through +// ggml-impl.h too, which includes the real header relative to ggml/src/), then +// layers on ONLY the block types ik_llama.cpp adds (iqk_common_extra.h). +// +// NO #pragma once: the upstream header is intentionally re-includable with +// different GGML_COMMON_DECL_*/IMPL_* each pass. The delta carries its own guard. +#include "../ggml-common.h" +#include "iqk_ggml_type_ext.h" // ik_llama ggml_type enum delta +#include "iqk_common_extra.h" diff --git a/llama.cpp/ggml/src/iqk/iqk_common.h b/llama.cpp/ggml/src/iqk/iqk_common.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_common.h @@ -0,0 +1,958 @@ +// -*- mode:c++;indent-tabs-mode:nil;c-basic-offset:4;coding:utf-8 -*- +// vi: set et ft=cpp fenc=utf-8 :vi +// +// +// Copyright (C) 2024 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include "iqk_config.h" + +#if defined IQK_IMPLEMENT + +#include +#include +#include +#include + +#include "ggml-impl.h" +#include "ggml-quants.h" +#include "iqk_mul_mat.h" +#include "iqk_quantize.h" + +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" + +#define FA_TIMING 0 + +#include +#include +#if FA_TIMING +#include +#include +struct Perf { + using TimePoint = std::chrono::time_point; + std::array times = {}; + std::mutex mutex; + bool report; + static auto cur_time() { return std::chrono::high_resolution_clock::now(); } + inline void accum(int what, const TimePoint& t1) { + auto t2 = cur_time(); + auto dt = delta(t1, t2); + std::lock_guard lock(mutex); + times[what] += dt; + } + inline void accum_nolock(int what, const TimePoint& t1) { + auto t2 = cur_time(); + auto dt = delta(t1, t2); + times[what] += dt; + } + inline void add(const Perf& other) { + std::lock_guard lock(mutex); + for (int i = 0; i < int(times.size()); ++i) times[i] += other.times[i]; + } + Perf(bool r) : report(r) {} + ~Perf() { + if (report) { + double tot = 0; + for (auto& t : times) tot += t; + if (!tot) return; + printf("======================= Timing: %g ms in total\n", tot); + for (int i = 0; i < int(times.size()); ++i) { + if (times[i]) { + printf("%d: %g ms -> %g%c\n", i, times[i], 100*times[i]/tot, '%'); + } + } + } + } + static Perf& instance() { + static Perf p(true); + return p; + } + static double delta(const TimePoint& t1, const TimePoint& t2) { + return 1e-6*std::chrono::duration_cast(t2-t1).count(); + } +}; +#endif + +#ifdef __AVX2__ +#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) +#endif + +typedef struct { + int32_t i1; + int32_t i2; +} mmid_row_mapping; + +struct DataInfo { + float * s; + const char * cy; + size_t bs; + size_t by; + int cur_y = 0; + int ne11; + const mmid_row_mapping * row_mapping = nullptr; + size_t bs2 = 0; + + inline const char * src1_row(int iy) const { + if (!row_mapping) return cy + (cur_y + iy)*by; + int i11 = row_mapping[cur_y + iy].i1 % ne11; + int i12 = row_mapping[cur_y + iy].i2; + return cy + (i11 + i12*ne11)*by; + } + + inline void store(int ix, int iy, float result) const { + *(dst_row(iy) + ix) = result; + } +#ifdef __AVX__ + inline void store(int ix, int iy, __m128 result) const { + _mm_storeu_ps(dst_row(iy) + ix, result); + } + inline void store(int ix, int iy, __m256 result) const { + _mm256_storeu_ps(dst_row(iy) + ix, result); + } +#endif +#ifdef __AVX512F__ + inline void store(int ix, int iy, __m512 result) const { + _mm512_storeu_ps(dst_row(iy) + ix, result); + } +#endif +#ifdef __ARM_NEON + inline void store(int ix, int iy, float32x4_t result) const { + vst1q_f32(dst_row(iy) + ix, result); + } +#endif + inline float * dst_row(int iy) const { + if (!row_mapping) return s + (cur_y + iy)*bs; + int i12 = row_mapping[cur_y + iy].i2; + int i1 = row_mapping[cur_y + iy].i1; + int i2 = i12; + return s + i1*bs + i2*bs2; + } +}; + +typedef void (*mul_mat_t)(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x); + +#define IQK_MAX_NY 8 + +#define IQK_SET_MUL_MAT_FUNCTIONS_T(kernel, Dequantizer, funcs) \ + funcs[0] = kernel;\ + funcs[1] = kernel;\ + funcs[2] = kernel;\ + funcs[3] = kernel;\ + funcs[4] = kernel;\ + funcs[5] = kernel;\ + funcs[6] = kernel;\ + funcs[7] = kernel;\ + +#define IQK_SET_MUL_MAT_FUNCTIONS_T2(kernel, Dequantizer, Block, funcs) \ + funcs[0] = kernel;\ + funcs[1] = kernel;\ + funcs[2] = kernel;\ + funcs[3] = kernel;\ + funcs[4] = kernel;\ + funcs[5] = kernel;\ + funcs[6] = kernel;\ + funcs[7] = kernel;\ + +#define IQK_SET_MUL_MAT_FUNCTIONS(kernel, funcs) \ + funcs[0] = kernel<1>;\ + funcs[1] = kernel<2>;\ + funcs[2] = kernel<3>;\ + funcs[3] = kernel<4>;\ + funcs[4] = kernel<5>;\ + funcs[5] = kernel<6>;\ + funcs[6] = kernel<7>;\ + funcs[7] = kernel<8>;\ + + +// ================================================================================================== + +static inline void make_q4_scales(const uint8_t * scales8, uint32_t * aux32) { + const uint16_t * scales = (const uint16_t *)scales8; + const uint32_t a0 = scales[0] | (scales[1] << 16); + const uint32_t a1 = scales[2] | (scales[3] << 16); + const uint32_t a2 = scales[4] | (scales[5] << 16); + aux32[3] = ((a2 >> 4) & 0x0f0f0f0f) | ((a1 >> 2) & 0x30303030); + aux32[1] = ((a2 >> 0) & 0x0f0f0f0f) | ((a0 >> 2) & 0x30303030); + aux32[2] = a1 & 0x3f3f3f3f; + aux32[0] = a0 & 0x3f3f3f3f; +} + +const uint64_t keven_signs[128] = { + 0x0101010101010101, 0xff010101010101ff, 0xff0101010101ff01, 0x010101010101ffff, + 0xff01010101ff0101, 0x0101010101ff01ff, 0x0101010101ffff01, 0xff01010101ffffff, + 0xff010101ff010101, 0x01010101ff0101ff, 0x01010101ff01ff01, 0xff010101ff01ffff, + 0x01010101ffff0101, 0xff010101ffff01ff, 0xff010101ffffff01, 0x01010101ffffffff, + 0xff0101ff01010101, 0x010101ff010101ff, 0x010101ff0101ff01, 0xff0101ff0101ffff, + 0x010101ff01ff0101, 0xff0101ff01ff01ff, 0xff0101ff01ffff01, 0x010101ff01ffffff, + 0x010101ffff010101, 0xff0101ffff0101ff, 0xff0101ffff01ff01, 0x010101ffff01ffff, + 0xff0101ffffff0101, 0x010101ffffff01ff, 0x010101ffffffff01, 0xff0101ffffffffff, + 0xff01ff0101010101, 0x0101ff01010101ff, 0x0101ff010101ff01, 0xff01ff010101ffff, + 0x0101ff0101ff0101, 0xff01ff0101ff01ff, 0xff01ff0101ffff01, 0x0101ff0101ffffff, + 0x0101ff01ff010101, 0xff01ff01ff0101ff, 0xff01ff01ff01ff01, 0x0101ff01ff01ffff, + 0xff01ff01ffff0101, 0x0101ff01ffff01ff, 0x0101ff01ffffff01, 0xff01ff01ffffffff, + 0x0101ffff01010101, 0xff01ffff010101ff, 0xff01ffff0101ff01, 0x0101ffff0101ffff, + 0xff01ffff01ff0101, 0x0101ffff01ff01ff, 0x0101ffff01ffff01, 0xff01ffff01ffffff, + 0xff01ffffff010101, 0x0101ffffff0101ff, 0x0101ffffff01ff01, 0xff01ffffff01ffff, + 0x0101ffffffff0101, 0xff01ffffffff01ff, 0xff01ffffffffff01, 0x0101ffffffffffff, + 0xffff010101010101, 0x01ff0101010101ff, 0x01ff01010101ff01, 0xffff01010101ffff, + 0x01ff010101ff0101, 0xffff010101ff01ff, 0xffff010101ffff01, 0x01ff010101ffffff, + 0x01ff0101ff010101, 0xffff0101ff0101ff, 0xffff0101ff01ff01, 0x01ff0101ff01ffff, + 0xffff0101ffff0101, 0x01ff0101ffff01ff, 0x01ff0101ffffff01, 0xffff0101ffffffff, + 0x01ff01ff01010101, 0xffff01ff010101ff, 0xffff01ff0101ff01, 0x01ff01ff0101ffff, + 0xffff01ff01ff0101, 0x01ff01ff01ff01ff, 0x01ff01ff01ffff01, 0xffff01ff01ffffff, + 0xffff01ffff010101, 0x01ff01ffff0101ff, 0x01ff01ffff01ff01, 0xffff01ffff01ffff, + 0x01ff01ffffff0101, 0xffff01ffffff01ff, 0xffff01ffffffff01, 0x01ff01ffffffffff, + 0x01ffff0101010101, 0xffffff01010101ff, 0xffffff010101ff01, 0x01ffff010101ffff, + 0xffffff0101ff0101, 0x01ffff0101ff01ff, 0x01ffff0101ffff01, 0xffffff0101ffffff, + 0xffffff01ff010101, 0x01ffff01ff0101ff, 0x01ffff01ff01ff01, 0xffffff01ff01ffff, + 0x01ffff01ffff0101, 0xffffff01ffff01ff, 0xffffff01ffffff01, 0x01ffff01ffffffff, + 0xffffffff01010101, 0x01ffffff010101ff, 0x01ffffff0101ff01, 0xffffffff0101ffff, + 0x01ffffff01ff0101, 0xffffffff01ff01ff, 0xffffffff01ffff01, 0x01ffffff01ffffff, + 0x01ffffffff010101, 0xffffffffff0101ff, 0xffffffffff01ff01, 0x01ffffffff01ffff, + 0xffffffffffff0101, 0x01ffffffffff01ff, 0x01ffffffffffff01, 0xffffffffffffffff, +}; + +#ifdef __AVX2__ + +#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) + +static inline float hsum_float_4(__m128 x) { + x = _mm_add_ps(x, _mm_movehl_ps(x, x)); + x = _mm_add_ss(x, _mm_movehdup_ps(x)); + return _mm_cvtss_f32(x); +} +static inline float hsum_float_8(__m256 x) { + return hsum_float_4(_mm_add_ps(_mm256_castps256_ps128(x), _mm256_extractf128_ps(x, 1))); +} +static inline int hsum_i32_8(const __m256i a) { + const __m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(a), _mm256_extractf128_si256(a, 1)); + const __m128i hi64 = _mm_unpackhi_epi64(sum128, sum128); + const __m128i sum64 = _mm_add_epi32(hi64, sum128); + const __m128i hi32 = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1)); + return _mm_cvtsi128_si32(_mm_add_epi32(sum64, hi32)); +} +static inline float hmax_f32_8(__m256 x) { + __m128 max4 = _mm_max_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x)); + max4 = _mm_max_ps(max4, _mm_movehl_ps(max4, max4)); + max4 = _mm_max_ss(max4, _mm_movehdup_ps(max4)); + return _mm_cvtss_f32(max4); +} +static inline float hmax_float_8(__m256 x) { + __m128 max4 = _mm_max_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x)); + max4 = _mm_max_ps( max4, _mm_movehl_ps(max4, max4)); + max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4)); + return _mm_cvtss_f32(max4); +} +static inline float hmin_float_8(__m256 x) { + __m128 min4 = _mm_min_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x)); + min4 = _mm_min_ps( min4, _mm_movehl_ps(min4, min4)); + min4 = _mm_min_ss( min4, _mm_movehdup_ps( min4)); + return _mm_cvtss_f32(min4); +} + +static inline __m128 hsum_float_4x4(__m128 * accm) { + accm[0] = _mm_add_ps(_mm_unpacklo_ps(accm[0], accm[2]), _mm_unpackhi_ps(accm[0], accm[2])); + accm[1] = _mm_add_ps(_mm_unpacklo_ps(accm[1], accm[3]), _mm_unpackhi_ps(accm[1], accm[3])); + return _mm_add_ps(_mm_unpacklo_ps(accm[0], accm[1]), _mm_unpackhi_ps(accm[0], accm[1])); +} +static inline __m256 hsum_float_8x8(__m256 * accm) { + for (int i = 0; i < 4; ++i) { + accm[i] = _mm256_add_ps(_mm256_permute2f128_ps(accm[i], accm[i + 4], 0x20), _mm256_permute2f128_ps(accm[i], accm[i + 4], 0x31)); + //accm[i] = _mm256_set_m128(_mm_add_ps(_mm256_castps256_ps128(accm[i+4]), _mm256_extractf128_ps(accm[i+4], 1)), + // _mm_add_ps(_mm256_castps256_ps128(accm[i+0]), _mm256_extractf128_ps(accm[i+0], 1))); + } + for (int i = 0; i < 2; ++i) accm[i] = _mm256_add_ps(_mm256_unpacklo_ps(accm[i], accm[i + 2]), _mm256_unpackhi_ps(accm[i], accm[i + 2])); + return _mm256_add_ps(_mm256_unpacklo_ps(accm[0], accm[1]), _mm256_unpackhi_ps(accm[0], accm[1])); +} +static inline __m256 hsum_float_4x8(__m256 * accm) { + for (int i = 0; i < 2; ++i) accm[i] = _mm256_add_ps(_mm256_unpacklo_ps(accm[i], accm[i + 2]), _mm256_unpackhi_ps(accm[i], accm[i + 2])); + return _mm256_add_ps(_mm256_unpacklo_ps(accm[0], accm[1]), _mm256_unpackhi_ps(accm[0], accm[1])); +} + +static inline __m128i load_iq4nl_values_128() { + static const uint8_t kvalues_iq4nl[16] = {1, 24, 45, 63, 79, 93, 106, 118, 129, 141, 153, 166, 181, 197, 217, 241}; + return _mm_loadu_si128((const __m128i *)kvalues_iq4nl); +} + +static inline __m256i load_iq4nl_values_256() { + auto val128 = load_iq4nl_values_128(); + return MM256_SET_M128I(val128, val128); +} + +#ifdef HAVE_FANCY_SIMD +static inline __m512i load_iq4nl_values_512() { + auto val256 = load_iq4nl_values_256(); + return _mm512_inserti32x8(_mm512_castsi256_si512(val256), val256, 1); +} +#endif + +static inline __m128i load_iq4k_values_128() { + return _mm_loadu_si128((const __m128i *)iq4k_values); +} + +static inline __m256i load_iq4k_values_256() { + auto val128 = load_iq4k_values_128(); + return MM256_SET_M128I(val128, val128); +} + +template struct Q8 { + + constexpr static int nrc_y = nrc; + + Q8(const DataInfo& info) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const block_q8 *)info.src1_row(iy); + } + +#ifdef HAVE_FANCY_SIMD + inline __m512i load_quants64(int iy, int i, int j) const { return _mm512_loadu_si512((const __m512i*)y[iy][i].qs + j); } +#endif + inline __m256i load_quants(int iy, int i, int j) const { return _mm256_loadu_si256((const __m256i*)y[iy][i].qs + j); } + inline __m256i load_bsums(int iy, int i) const { return _mm256_loadu_si256((const __m256i*)y[iy][i].bsums); } + inline float scale(int iy, int i) const { return y[iy][i].d; } + + const block_q8 * y[nrc_y]; +}; + +template struct Q8_16 { + + constexpr static int nrc_y = nrc; + + Q8_16(const DataInfo& info) { + for (int iy = 0; iy < nrc_y; ++iy) { + auto ptr = (const float *)info.src1_row(iy); + std::memcpy(d + 5*iy, ptr, 5*sizeof(float)); + y[iy] = (const int8_t *)(ptr + 5); + } + } + +#ifdef HAVE_FANCY_SIMD + inline __m512i load_quants64(int iy, int i) const { return _mm512_loadu_si512((const __m512i*)y[iy] + i); } +#endif + inline __m256i load_quants(int iy, int i) const { return _mm256_loadu_si256((const __m256i*)y[iy] + i); } + inline float scale(int iy, int k) const { return d[5*iy+k]; } + inline float sum_row(int iy) const { return d[5*iy + 4]; } + inline __m128 scale(int iy) const { return _mm_loadu_ps(d + 5*iy); } + + float d[5*nrc_y]; + const int8_t * y[nrc_y]; +}; + +struct Scales8KBase { + template + inline void accum_mins(const __m128i& mins128, const Q8& q8, int i, float c, __m256 * accd) const { + const __m256i mins = MM256_SET_M128I(_mm_shuffle_epi8(mins128, shuffles[1]), _mm_shuffle_epi8(mins128, shuffles[0])); + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + const __m256i q8s = q8.load_bsums(iy, i); + const __m256i prod = _mm256_madd_epi16(mins, q8s); + accd[iy] = _mm256_fmadd_ps(_mm256_set1_ps(c*q8.scale(iy, i)), _mm256_cvtepi32_ps(prod), accd[iy]); + } + } + inline __m256i shuffle(__m128i mins) const { + return MM256_SET_M128I(_mm_shuffle_epi8(mins, shuffles[1]), _mm_shuffle_epi8(mins, shuffles[0])); + } + const __m128i shuffles[2] = {_mm_set_epi32(0x07060706, 0x05040504, 0x03020302, 0x01000100), + _mm_set_epi32(0x0f0e0f0e, 0x0d0c0d0c, 0x0b0a0b0a, 0x09080908)}; +}; + +template +struct BaseDequantizer { + BaseDequantizer(const void * vx, size_t bx) : vx(vx), bx(bx) {} + inline void new_row(int ix) { + if constexpr (per_row_scale) { + if constexpr (is_f16) { + const ggml_half * dptr = (const ggml_half *)((const char *)vx + bx*ix); + d = GGML_FP16_TO_FP32(*dptr); + x = (const Block *)(dptr + 1); + } else { + const float * dptr = (const float *)((const char *)vx + bx*ix); + d = *dptr; + x = (const Block *)(dptr + 1); + } + } else { + x = (const Block *)((const char *)vx + bx*ix); + } + } + + const void * vx; + const size_t bx; + const Block * x; + + float d; +}; + +template +static inline void multiply_add(const Bits& bits, const __m256i * scales, int j, int i, const Q8& q8, __m256i * sumi) { + if (j == 0) { +#ifdef HAVE_FANCY_SIMD + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + sumi[iy] = _mm256_dpwssd_epi32(_mm256_setzero_si256(), scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 0))); + sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 1))); + sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 2))); + sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 3))); + } +#else + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + const __m256i p1 = _mm256_madd_epi16(scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 0))); + const __m256i p2 = _mm256_madd_epi16(scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 1))); + const __m256i p3 = _mm256_madd_epi16(scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 2))); + const __m256i p4 = _mm256_madd_epi16(scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 3))); + sumi[iy] = _mm256_add_epi32(_mm256_add_epi32(p1, p3), _mm256_add_epi32(p2, p4)); + } +#endif + } else { +#ifdef HAVE_FANCY_SIMD + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 4))); + sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 5))); + sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 6))); + sumi[iy] = _mm256_dpwssd_epi32(sumi[iy], scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 7))); + } +#else + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + const __m256i p1 = _mm256_madd_epi16(scales[0], _mm256_maddubs_epi16(bits.values[0], q8.load_quants(iy, i, 4))); + const __m256i p2 = _mm256_madd_epi16(scales[1], _mm256_maddubs_epi16(bits.values[1], q8.load_quants(iy, i, 5))); + const __m256i p3 = _mm256_madd_epi16(scales[2], _mm256_maddubs_epi16(bits.values[2], q8.load_quants(iy, i, 6))); + const __m256i p4 = _mm256_madd_epi16(scales[3], _mm256_maddubs_epi16(bits.values[3], q8.load_quants(iy, i, 7))); + sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p1, p3)); + sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p2, p4)); + } +#endif + } +} + +template +static inline void multiply_add_avx2(const Bits& bits, const __m256i * scales, int j, int i, const Q8& q8, __m256i * sumi) { + __m256i p[4]; + if (j == 0) { + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + for (int k = 0; k < 4; ++k) { + auto s = _mm256_sign_epi8(bits.values[k], bits.values[k]); + p[k] = _mm256_madd_epi16(scales[k], _mm256_maddubs_epi16(s, _mm256_sign_epi8(q8.load_quants(iy, i, k), bits.values[k]))); + } + sumi[iy] = _mm256_add_epi32(_mm256_add_epi32(p[0], p[1]), _mm256_add_epi32(p[2], p[3])); + } + } else { + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + for (int k = 0; k < 4; ++k) { + auto s = _mm256_sign_epi8(bits.values[k], bits.values[k]); + p[k] = _mm256_madd_epi16(scales[k], _mm256_maddubs_epi16(s, _mm256_sign_epi8(q8.load_quants(iy, i, 4+k), bits.values[k]))); + } + sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p[0], p[2])); + sumi[iy] = _mm256_add_epi32(sumi[iy], _mm256_add_epi32(p[1], p[3])); + } + } +} + +#ifdef HAVE_FANCY_SIMD + +struct BlockPermuter { + const __m512i permute1 = _mm512_set_epi64(11, 10, 9, 8, 3, 2, 1, 0); + const __m512i permute2 = _mm512_set_epi64(15, 14, 13, 12, 7, 6, 5, 4); +}; + +struct Q4Bits { + inline void prepare(const uint8_t * q4) { + auto q4bits = _mm512_loadu_si512((const __m512i*)q4 + 0); + auto tmp1 = _mm512_and_si512(q4bits, ml); + auto tmp2 = _mm512_and_si512(_mm512_srli_epi16(q4bits, 4), ml); + values[0] = _mm512_permutex2var_epi64(tmp1, perm.permute1, tmp2); + values[1] = _mm512_permutex2var_epi64(tmp1, perm.permute2, tmp2); + q4bits = _mm512_loadu_si512((const __m512i*)q4 + 1); + tmp1 = _mm512_and_si512(q4bits, ml); + tmp2 = _mm512_and_si512(_mm512_srli_epi16(q4bits, 4), ml); + values[2] = _mm512_permutex2var_epi64(tmp1, perm.permute1, tmp2); + values[3] = _mm512_permutex2var_epi64(tmp1, perm.permute2, tmp2); + } + inline void prepare64(const uint8_t * q4) { + auto q4bits = _mm512_loadu_si512((const __m512i*)q4 + 0); + values[0] = _mm512_and_si512(q4bits, ml); + values[1] = _mm512_and_si512(_mm512_srli_epi16(q4bits, 4), ml); + q4bits = _mm512_loadu_si512((const __m512i*)q4 + 1); + values[2] = _mm512_and_si512(q4bits, ml); + values[3] = _mm512_and_si512(_mm512_srli_epi16(q4bits, 4), ml); + } + inline void prepare64a(const uint8_t * q4) { + for (int k = 0; k < 4; ++k) { + auto q4bits = _mm256_loadu_si256((const __m256i*)q4 + k); + values[k] = _mm512_inserti32x8(_mm512_castsi256_si512(q4bits), _mm256_srli_epi16(q4bits, 4), 1); + values[k] = _mm512_and_si512(values[k], ml); + } + } + __m512i values[4]; + const __m512i ml = _mm512_set1_epi8(0xf); + const BlockPermuter perm; +}; + +struct Q2Bits { + inline void prepare(const uint8_t * q2) { + + auto q2bits = _mm512_loadu_si512((const __m512i*)q2); + auto tmp = _mm512_srli_epi16(q2bits, 2); + + values[0] = _mm512_permutex2var_epi64(q2bits, perm.permute1, tmp); + values[2] = _mm512_permutex2var_epi64(q2bits, perm.permute2, tmp); + values[1] = _mm512_and_si512(_mm512_srli_epi16(values[0], 4), ml); + values[3] = _mm512_and_si512(_mm512_srli_epi16(values[2], 4), ml); + values[0] = _mm512_and_si512(values[0], ml); + values[2] = _mm512_and_si512(values[2], ml); + } + __m512i values[4]; + const __m512i ml = _mm512_set1_epi8(0x03); + BlockPermuter perm; +}; + +#else + +struct Q2Bits { + inline void prepare(const uint8_t * q2, int j) { + auto q2bits = _mm256_loadu_si256((const __m256i *)q2 + j); + values[0] = _mm256_and_si256(q2bits, ml); + values[1] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 2), ml); + values[2] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 4), ml); + values[3] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 6), ml); + } + __m256i values[4]; + const __m256i ml = _mm256_set1_epi8(0x03); +}; + +struct Q4Bits { + inline void prepare(const uint8_t * q4, int j) { + auto q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+0); + values[0] = _mm256_and_si256(q4bits, ml); + values[1] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); + q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+1); + values[2] = _mm256_and_si256(q4bits, ml); + values[3] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); + } + inline void prepare64(const uint8_t * q4, int j) { + auto q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+0); + values[0] = _mm256_and_si256(q4bits, ml); + values[2] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); + q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+1); + values[1] = _mm256_and_si256(q4bits, ml); + values[3] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); + } + inline void prepare16(const uint8_t * q4, int j) { + values[0] = dequant16(q4 + 64*j + 0); + values[1] = dequant16(q4 + 64*j + 16); + values[2] = dequant16(q4 + 64*j + 32); + values[3] = dequant16(q4 + 64*j + 48); + } + inline __m256i dequant16(const uint8_t * qs) const { + const __m128i aux128 = _mm_loadu_si128((const __m128i *)qs); + const __m256i aux256 = MM256_SET_M128I(_mm_srli_epi16(aux128, 4), aux128); + return _mm256_and_si256(ml, aux256); + } + __m256i values[4]; + const __m256i ml = _mm256_set1_epi8(0xf); +}; + +#endif + +inline void iqk_transpose_8x8(__m256 * m) { + for (int k = 0; k < 8; k += 4) { + auto t0 = _mm256_unpacklo_ps(m[k+0], m[k+1]); + auto t1 = _mm256_unpacklo_ps(m[k+2], m[k+3]); + auto t2 = _mm256_unpackhi_ps(m[k+0], m[k+1]); + auto t3 = _mm256_unpackhi_ps(m[k+2], m[k+3]); + m[k+0] = _mm256_castpd_ps(_mm256_unpacklo_pd(_mm256_castps_pd(t0), _mm256_castps_pd(t1))); + m[k+1] = _mm256_castpd_ps(_mm256_unpackhi_pd(_mm256_castps_pd(t0), _mm256_castps_pd(t1))); + m[k+2] = _mm256_castpd_ps(_mm256_unpacklo_pd(_mm256_castps_pd(t2), _mm256_castps_pd(t3))); + m[k+3] = _mm256_castpd_ps(_mm256_unpackhi_pd(_mm256_castps_pd(t2), _mm256_castps_pd(t3))); + } + for (int k = 0; k < 4; ++k) { + auto t = _mm256_set_m128(_mm256_extractf128_ps(m[k+4], 1), _mm256_extractf128_ps(m[k], 1)); + m[k+0] = _mm256_set_m128(_mm256_castps256_ps128(m[k+4]), _mm256_castps256_ps128(m[k+0])); + m[k+4] = t; + } +} + +template +static inline float convert_to_q8_k_r8(int k, float d0, const __m256i * qx, const int16_t * scales, uint32_t * block, int8_t * q8_k) { + auto max_i16 = _mm256_setzero_si256(); + __m256i qs[16]; + for (int ib32 = 0; ib32 < 8; ++ib32) { + qs[2*ib32+0] = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(qx[ib32])); + qs[2*ib32+1] = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(qx[ib32], 1)); + qs[2*ib32+0] = _mm256_mullo_epi16(qs[2*ib32+0], _mm256_set1_epi16(scales[2*ib32+0])); + qs[2*ib32+1] = _mm256_mullo_epi16(qs[2*ib32+1], _mm256_set1_epi16(scales[2*ib32+1])); + max_i16 = _mm256_max_epi16(max_i16, _mm256_sign_epi16(qs[2*ib32+0], qs[2*ib32+0])); + max_i16 = _mm256_max_epi16(max_i16, _mm256_sign_epi16(qs[2*ib32+1], qs[2*ib32+1])); + } + auto max_q32 = _mm256_cvtepi16_epi32(_mm_max_epi16(_mm256_castsi256_si128(max_i16), _mm256_extracti128_si256(max_i16, 1))); + auto imax4 = _mm_max_epi32(_mm256_castsi256_si128(max_q32), _mm256_extracti128_si256(max_q32, 1)); + auto max4 = _mm_cvtepi32_ps(imax4); + max4 = _mm_max_ps(max4, _mm_movehl_ps(max4, max4)); + max4 = _mm_max_ss(max4, _mm_movehdup_ps(max4)); + bool needs_scaling = true; + float dnew = _mm_cvtss_f32(max4) * d0; + if (dnew < 1.f) { + dnew = 1.f; needs_scaling = false; + } + auto scale = _mm256_set1_ps(std::abs(dnew) > 1e-9f ? 1/dnew : 0.f); + for (int ib32 = 0; ib32 < 8; ++ib32) { + if (needs_scaling) { + auto i0 = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(qs[2*ib32+0])); + auto i1 = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(qs[2*ib32+0], 1)); + auto i2 = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(qs[2*ib32+1])); + auto i3 = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(qs[2*ib32+1], 1)); + i0 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i0)), _MM_ROUND_NEAREST)); + i1 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i1)), _MM_ROUND_NEAREST)); + i2 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i2)), _MM_ROUND_NEAREST)); + i3 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i3)), _MM_ROUND_NEAREST)); + i0 = _mm256_packs_epi32(i0, i1); + i2 = _mm256_packs_epi32(i2, i3); + i0 = _mm256_packs_epi16(i0, i2); + i0 = _mm256_permutevar8x32_epi32(i0, _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7)); + _mm256_storeu_si256((__m256i *)block, i0); + } else { + // 0, 1, 2, 3, 4, 5, 6, 7, 8, 16, 17, 18, 19, 20, 21, 22, 23, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31 + auto i0 = _mm256_packs_epi16(qs[2*ib32+0], qs[2*ib32+1]); + auto i0_l = _mm256_castsi256_si128(i0); + auto i0_h = _mm256_extracti128_si256(i0, 1); + _mm_storeu_si128((__m128i *)block+0, _mm_unpacklo_epi64(i0_l, i0_h)); + _mm_storeu_si128((__m128i *)block+1, _mm_unpackhi_epi64(i0_l, i0_h)); + } + auto qs = (uint32_t *)q8_k + 8*nr*ib32; + for (int l = 0; l < 8; ++l) { + qs[nr*l + k] = block[l]; + } + } + return dnew; +} + +#else +// ------------------------------------ __aarch64__ -------------------------------------------------- + +template struct Q8 { + + constexpr static int nrc_y = nrc; + + Q8(const DataInfo& info) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const block_q8 *)info.src1_row(iy); + } + + inline int8x16x2_t load_quants(int iy, int i, int j) const { return vld1q_s8_x2(y[iy][i].qs + 32*j); } + inline int8x16x4_t load_quants_64(int iy, int i, int j) const { return vld1q_s8_x4(y[iy][i].qs + 64*j); } + inline int16x8x2_t load_bsums(int iy, int i) const { return vld1q_s16_x2(y[iy][i].bsums); } + inline int16x8_t load_bsums8(int iy, int i) const { + auto q8s = vld1q_s16_x2(y[iy][i].bsums); + return vpaddq_s16(q8s.val[0], q8s.val[1]); + } + inline float scale(int iy, int i) const { return y[iy][i].d; } + + const block_q8 * y[nrc_y]; +}; + +template +struct BaseDequantizer { + BaseDequantizer(const void * vx, size_t bx, int nrc) : vx(vx), x(nullptr), bx(bx), nrc(nrc) {} + inline void new_row(int ix) { + if constexpr (has_row_scale) { + if constexpr (scale_is_f16) { + const ggml_half * dptr = (const ggml_half *)((const char *)vx + ix*bx); + d = GGML_FP16_TO_FP32(*dptr); + x = (const block_q *)(dptr + 1); + } else { + const float * dptr = (const float *)((const char *)vx + ix*bx); + d = *dptr; + x = (const block_q *)(dptr + 1); + } + } else { + x = (const block_q *)((const char *)vx + ix*bx); + } + } + const void * vx; + const block_q * x; + const size_t bx; + const int nrc; + float d; +}; + +struct Q4bits { + const uint8x16_t m4b = vdupq_n_u8(0xf); + uint8x16x4_t b1, b2; + inline void prepare4(uint8x16x4_t& b, const uint8x16_t * val) const { + b.val[0] = vandq_u8(val[0], m4b); + b.val[2] = vshrq_n_u8(val[0], 4); + b.val[1] = vandq_u8(val[1], m4b); + b.val[3] = vshrq_n_u8(val[1], 4); + } + inline void prepare4_16(uint8x16x4_t& b, const uint8x16_t * val) const { + b.val[0] = vandq_u8(val[0], m4b); + b.val[1] = vshrq_n_u8(val[0], 4); + b.val[2] = vandq_u8(val[1], m4b); + b.val[3] = vshrq_n_u8(val[1], 4); + } + inline void prepare(const uint8_t * qs) { + auto q4bits = vld1q_u8_x2(qs); + prepare4(b1, q4bits.val); + q4bits = vld1q_u8_x2(qs+32); + prepare4(b2, q4bits.val); + } + inline void prepare_v2(const uint8_t * qs) { + auto q4bits = vld1q_u8_x4(qs); + prepare4(b1, q4bits.val+0); + prepare4(b2, q4bits.val+2); + } + inline void prepare64(const uint8_t * qs) { + auto q4bits = vld1q_u8_x4(qs); + b1.val[0] = vandq_u8(q4bits.val[0], m4b); + b1.val[1] = vandq_u8(q4bits.val[1], m4b); + b1.val[2] = vandq_u8(q4bits.val[2], m4b); + b1.val[3] = vandq_u8(q4bits.val[3], m4b); + b2.val[0] = vshrq_n_u8(q4bits.val[0], 4); + b2.val[1] = vshrq_n_u8(q4bits.val[1], 4); + b2.val[2] = vshrq_n_u8(q4bits.val[2], 4); + b2.val[3] = vshrq_n_u8(q4bits.val[3], 4); + } + inline void prepare16(const uint8_t * qs) { + auto q4bits = vld1q_u8_x2(qs); + prepare4_16(b1, q4bits.val); + q4bits = vld1q_u8_x2(qs+32); + prepare4_16(b2, q4bits.val); + } + inline void prepare16_v2(const uint8_t * qs) { + auto q4bits = vld1q_u8_x4(qs); + prepare4_16(b1, q4bits.val+0); + prepare4_16(b2, q4bits.val+2); + } +}; + +struct Q2bits { + const uint8x16_t m4b = vdupq_n_u8(0x03); + uint8x16x4_t b1, b2; + inline void prepare(const uint8_t * qs) { + auto q2bits = vld1q_u8_x2(qs); + b1.val[0] = vandq_u8(q2bits.val[0], m4b); + b1.val[1] = vandq_u8(q2bits.val[1], m4b); + + q2bits.val[0] = vshrq_n_u8(q2bits.val[0], 2); + q2bits.val[1] = vshrq_n_u8(q2bits.val[1], 2); + b1.val[2] = vandq_u8(q2bits.val[0], m4b); + b1.val[3] = vandq_u8(q2bits.val[1], m4b); + + q2bits.val[0] = vshrq_n_u8(q2bits.val[0], 2); + q2bits.val[1] = vshrq_n_u8(q2bits.val[1], 2); + b2.val[0] = vandq_u8(q2bits.val[0], m4b); + b2.val[1] = vandq_u8(q2bits.val[1], m4b); + + q2bits.val[0] = vshrq_n_u8(q2bits.val[0], 2); + q2bits.val[1] = vshrq_n_u8(q2bits.val[1], 2); + b2.val[2] = vandq_u8(q2bits.val[0], m4b); + b2.val[3] = vandq_u8(q2bits.val[1], m4b); + } +}; + +template +static inline void compute_8_blocks(const uint8x16x4_t& qx_1, const uint8x16x4_t& qx_2, const Q8& q8, + const int32x4x2_t& scales, int iy, int i, int j, int32x4_t& sumi) { + auto mzero = vdupq_n_s32(0); + auto q8b_1 = q8.load_quants(iy, i, 4*j+0); + auto p1 = ggml_vdotq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[0]), q8b_1.val[0]), + vreinterpretq_s8_u8(qx_1.val[1]), q8b_1.val[1]); // block 1 + auto q8b_2 = q8.load_quants(iy, i, 4*j+1); + auto p2 = ggml_vdotq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[2]), q8b_2.val[0]), + vreinterpretq_s8_u8(qx_1.val[3]), q8b_2.val[1]); // block 2 + auto p12 = vpaddq_s32(p1, p2); + + auto q8b_3 = q8.load_quants(iy, i, 4*j+2); + auto p3 = ggml_vdotq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[0]), q8b_3.val[0]), + vreinterpretq_s8_u8(qx_2.val[1]), q8b_3.val[1]); // block 1 + auto q8b_4 = q8.load_quants(iy, i, 4*j+3); + auto p4 = ggml_vdotq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[2]), q8b_4.val[0]), + vreinterpretq_s8_u8(qx_2.val[3]), q8b_4.val[1]); // block 2 + auto p34 = vpaddq_s32(p3, p4); + + auto pall = vpaddq_s32(p12, p34); + sumi = vmlaq_s32(sumi, scales.val[j], pall); +} + +template +static inline void compute_16_blocks(const uint8x16x4_t& qx_1, const uint8x16x4_t& qx_2, const Q8& q8, + const int32x4x4_t& scales, int iy, int i, int j, int32x4_t& sumi) { + + auto mzero = vdupq_n_s32(0); + auto q8b_1 = q8.load_quants(iy, i, 4*j+0); + auto p1 = vpaddq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[0]), q8b_1.val[0]), + ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[1]), q8b_1.val[1])); // blocks 0, 0, 1, 1, + auto q8b_2 = q8.load_quants(iy, i, 4*j+1); + auto p2 = vpaddq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[2]), q8b_2.val[0]), + ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_1.val[3]), q8b_2.val[1])); // blocks 3, 3, 4, 4, + auto p12 = vpaddq_s32(p1, p2); // blocks 0, 1, 2, 3 + sumi = vmlaq_s32(sumi, scales.val[2*j+0], p12); + + auto q8b_3 = q8.load_quants(iy, i, 4*j+2); + auto p3 = vpaddq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[0]), q8b_3.val[0]), + ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[1]), q8b_3.val[1])); // block 4, 4, 5, 5, + auto q8b_4 = q8.load_quants(iy, i, 4*j+3); + auto p4 = vpaddq_s32(ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[2]), q8b_4.val[0]), + ggml_vdotq_s32(mzero, vreinterpretq_s8_u8(qx_2.val[3]), q8b_4.val[1])); // block 6, 6, 7, 7, + auto p34 = vpaddq_s32(p3, p4); // blocks 4, 5, 6, 7 + sumi = vmlaq_s32(sumi, scales.val[2*j+1], p34); +} + +struct SignHelper { + + inline void init() { shuffle = vcombine_u8(vdup_n_u8(0), vdup_n_u8(1)); } + + inline void apply_signs_1(uint8x16_t * b, const uint8x16_t& signs16) { + auto aux = vqtbl1q_u8(signs16, shuffle); + auto s = vreinterpretq_s8_u8(vorrq_u8(vceqq_u8(vandq_u8(aux, smask), smask), m1)); + b[0] = vreinterpretq_u8_s8(vmulq_s8(vreinterpretq_s8_u8(b[0]), s)); + shuffle = vaddq_u8(shuffle, step); + } + + const uint8x16_t smask = vreinterpretq_u8_u64(vdupq_n_u64(0x8040201008040201)); + const uint8x16_t m1 = vdupq_n_u8(1); + const uint8x16_t step = vdupq_n_u8(2); + uint8x16_t shuffle; +}; + +template +static void mul_mat_qX_K_q8_K_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + Q8 q8(info); + + Dequantizer deq(vx, bx, nrc_y); + + for (int ix = 0; ix < nrc_x; ++ix) { + + deq.new_row(ix); + + float32x4_t acc[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) acc[iy] = vdupq_n_f32(0.f); + + for (int i = 0; i < nb; ++i) { + + int32x4_t sumi[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) sumi[iy] = vdupq_n_s32(0); + + if constexpr (nrc_y > 1 && Dequantizer::should_scale_quants()) { + deq.process_scales(i, q8, acc); + deq.prepare(i, 0); + deq.compute(q8, i, 0, sumi); + deq.prepare(i, 1); + deq.compute(q8, i, 1, sumi); + } else { + if constexpr (Dequantizer::num_blocks() == 8) { + auto scales = deq.new_block(i, q8, acc); + deq.prepare(i, 0); + for (int iy = 0; iy < nrc_y; ++iy) compute_8_blocks(deq.bits.b1, deq.bits.b2, q8, scales, iy, i, 0, sumi[iy]); + deq.prepare(i, 1); + for (int iy = 0; iy < nrc_y; ++iy) compute_8_blocks(deq.bits.b1, deq.bits.b2, q8, scales, iy, i, 1, sumi[iy]); + } + else if constexpr (Dequantizer::num_blocks() == 16) { + auto scales = deq.new_block(i, q8, acc); + deq.prepare(i, 0); + for (int iy = 0; iy < nrc_y; ++iy) compute_16_blocks(deq.bits.b1, deq.bits.b2, q8, scales, iy, i, 0, sumi[iy]); + deq.prepare(i, 1); + for (int iy = 0; iy < nrc_y; ++iy) compute_16_blocks(deq.bits.b1, deq.bits.b2, q8, scales, iy, i, 1, sumi[iy]); + } + else { + GGML_ASSERT(false); + } + } + + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = vmlaq_f32(acc[iy], vcvtq_f32_s32(sumi[iy]), vdupq_n_f32(deq.d*q8.scale(iy, i))); + } + } + + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, vaddvq_f32(acc[iy])); + } + } +} + +static IQK_ALWAYS_INLINE int32x4_t interleaved_dotq(const int8x16_t * qx, const int8x16x2_t& y) { + auto sumi = vdupq_n_s32(0); + sumi = vdotq_laneq_s32(sumi, qx[0], y.val[0], 0); + sumi = vdotq_laneq_s32(sumi, qx[1], y.val[1], 0); + sumi = vdotq_laneq_s32(sumi, qx[2], y.val[0], 1); + sumi = vdotq_laneq_s32(sumi, qx[3], y.val[1], 1); + sumi = vdotq_laneq_s32(sumi, qx[4], y.val[0], 2); + sumi = vdotq_laneq_s32(sumi, qx[5], y.val[1], 2); + sumi = vdotq_laneq_s32(sumi, qx[6], y.val[0], 3); + sumi = vdotq_laneq_s32(sumi, qx[7], y.val[1], 3); + return sumi; +} + +static IQK_ALWAYS_INLINE int32x4x2_t interleaved_dotq_b16(const int8x16_t * qx, const int8x16x2_t& y) { + int32x4x2_t sumi = { vdupq_n_s32(0), vdupq_n_s32(0) }; + sumi.val[0] = vdotq_laneq_s32(sumi.val[0], qx[0], y.val[0], 0); + sumi.val[1] = vdotq_laneq_s32(sumi.val[1], qx[1], y.val[1], 0); + sumi.val[0] = vdotq_laneq_s32(sumi.val[0], qx[2], y.val[0], 1); + sumi.val[1] = vdotq_laneq_s32(sumi.val[1], qx[3], y.val[1], 1); + sumi.val[0] = vdotq_laneq_s32(sumi.val[0], qx[4], y.val[0], 2); + sumi.val[1] = vdotq_laneq_s32(sumi.val[1], qx[5], y.val[1], 2); + sumi.val[0] = vdotq_laneq_s32(sumi.val[0], qx[6], y.val[0], 3); + sumi.val[1] = vdotq_laneq_s32(sumi.val[1], qx[7], y.val[1], 3); + return sumi; +} + +static IQK_ALWAYS_INLINE int32x4_t interleaved_dotq(const int8x16_t * qx, const int8x16_t& y) { + auto sumi = vdupq_n_s32(0); + sumi = vdotq_laneq_s32(sumi, qx[0], y, 0); + sumi = vdotq_laneq_s32(sumi, qx[1], y, 1); + sumi = vdotq_laneq_s32(sumi, qx[2], y, 2); + sumi = vdotq_laneq_s32(sumi, qx[3], y, 3); + return sumi; +} + +static IQK_ALWAYS_INLINE void prepare_iq4_nl_quants(const int8x16_t& values, const uint8x16_t& m4, const uint8x16x4_t& bits, int8x16_t * qx) { + qx[0] = vqtbl1q_s8(values, vandq_u8(bits.val[0], m4)); // 0...3 from the 4 rows + qx[1] = vqtbl1q_s8(values, vandq_u8(bits.val[1], m4)); // 16..19 + qx[2] = vqtbl1q_s8(values, vandq_u8(bits.val[2], m4)); // 4...7 + qx[3] = vqtbl1q_s8(values, vandq_u8(bits.val[3], m4)); // 20..23 + qx[4] = vqtbl1q_s8(values, vshrq_n_u8(bits.val[0], 4)); // 8..11 + qx[5] = vqtbl1q_s8(values, vshrq_n_u8(bits.val[1], 4)); // 24..27 + qx[6] = vqtbl1q_s8(values, vshrq_n_u8(bits.val[2], 4)); // 12..15 + qx[7] = vqtbl1q_s8(values, vshrq_n_u8(bits.val[3], 4)); // 28..31 +} + +static IQK_ALWAYS_INLINE void prepare_iq4_nl_quants_r8(const int8x16_t& values, const uint8x16_t& m4, const uint8x16x2_t& bits, int8x16_t * qx) { + qx[0] = vqtbl1q_s8(values, vandq_u8( bits.val[0], m4)); + qx[1] = vqtbl1q_s8(values, vshrq_n_u8(bits.val[0], 4)); + qx[2] = vqtbl1q_s8(values, vandq_u8( bits.val[1], m4)); + qx[3] = vqtbl1q_s8(values, vshrq_n_u8(bits.val[1], 4)); +} + +#endif + +#endif + +// static unrool for: +template +inline void static_for(T&&f) { + if constexpr(N>0) { + static_for(f); + f(N-1); + } +} + +#if defined(_MSC_VER) +#pragma warning(disable: 4244 4267) // possible loss of data +#include +#include +#include +#include +#include +inline int popcount(uint8_t x) { return __popcnt(x); } +inline int popcount(uint16_t x) { return __popcnt(x); } +inline int popcount(uint32_t x) { return __popcnt(x); } +inline int popcount(uint64_t x) { return _mm_popcnt_u64(x); } +#else +constexpr int popcount(uint8_t x) { return __builtin_popcount(x); } +constexpr int popcount(uint16_t x) { return __builtin_popcount(x); } +constexpr int popcount(uint32_t x) { return __builtin_popcount(x); } +constexpr int popcount(uint64_t x) { return __builtin_popcountll(x); } +#endif + diff --git a/llama.cpp/ggml/src/iqk/iqk_common_extra.h b/llama.cpp/ggml/src/iqk/iqk_common_extra.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_common_extra.h @@ -0,0 +1,471 @@ +// opencoti F5-opt W2 (#290) — ik_llama.cpp ggml-common.h DELTA. +// VENDORED from ik_llama.cpp @ 8960c5ba… ggml/src/ggml-common.h +// Contains ONLY the block types ik_llama defines that llamafile's ggml-common.h +// lacks (extracted compiler-driven: each entry was an "undeclared type" error). +// MIT, (C) 2024 Iwan Kawrakow. Re-sync: diff upstream ggml-common.h block defs. +#pragma once +#ifndef IQK_COMMON_EXTRA_H +#define IQK_COMMON_EXTRA_H + +typedef struct { + ggml_half d[4]; // delta + uint8_t qh[QK5_0/2]; // 5-th bit of quants + uint8_t qs[QK5_0*2]; // nibbles / quants +} block_q5_0_r4; +static_assert(sizeof(block_q5_0_r4) == 4*sizeof(ggml_half) + QK5_0*2 + QK5_0/2, "wrong q5_0_r4 block size/padding"); + +#define QK6_0 32 +typedef struct { + ggml_half d; // delta + uint8_t qh[QK6_0/4]; // 5+6-th bit of quants + uint8_t qs[QK6_0/2]; // nibbles / quants +} block_q6_0; +static_assert(sizeof(block_q6_0) == sizeof(ggml_half) + QK6_0/2 + QK6_0/4, "wrong q6_0 block size/padding"); + +typedef struct { + ggml_half d[4]; // delta + uint8_t qh[QK6_0]; // 5+6-th bit of quants + uint8_t qs[QK6_0*2]; // nibbles / quants +} block_q6_0_r4; +static_assert(sizeof(block_q6_0_r4) == 4*sizeof(ggml_half) + QK6_0*2 + QK6_0, "wrong q6_0_r4 block size/padding"); + +typedef struct { + ggml_half d[8]; + int8_t qs[4*QK8_1]; +} block_q8_1_x4; +static_assert(sizeof(block_q8_1_x4) == 4*sizeof(block_q8_1), "wrong q8_1_x4 block size/padding"); + +typedef struct { + ggml_half d[4]; + int8_t qs[4*QK8_0]; +} block_q8_0_x4; +static_assert(sizeof(block_q8_0_x4) == 4*sizeof(block_q8_0), "wrong q8_0_x4 block size/padding"); + +typedef struct { + ggml_half d[8]; + int8_t qs[8*QK8_0]; +} block_q8_0_r8; +static_assert(sizeof(block_q8_0_r8) == 8*sizeof(block_q8_0), "wrong q8_0_r8 block size/padding"); + +typedef struct { + ggml_half d[4]; // deltas for 4 q4_0 blocks + uint8_t qs[QK4_0 * 2]; // nibbles / quants for 4 q4_0 blocks +} block_q4_0x4; +static_assert(sizeof(block_q4_0x4) == 4 * sizeof(ggml_half) + QK4_0 * 2, "wrong q4_0x4 block size/padding"); + +typedef struct { + ggml_half d[8]; // deltas for 8 q4_0 blocks + uint8_t qs[QK4_0 * 4]; // nibbles / quants for 8 q4_0 blocks +} block_q4_0x8; +static_assert(sizeof(block_q4_0x8) == 8 * sizeof(ggml_half) + QK4_0 * 4, "wrong q4_0x8 block size/padding"); + +typedef struct { + ggml_half d[4]; // deltas for 4 q8_0 blocks + int8_t qs[QK8_0 * 4]; // quants for 4 q8_0 blocks +} block_q8_0x4; +static_assert(sizeof(block_q8_0x4) == 4 * sizeof(ggml_half) + QK8_0 * 4, "wrong q8_0x4 block size/padding"); + +typedef struct { + ggml_half d[8]; // deltas for 8 q8_0 blocks + int8_t qs[QK8_0 * 8]; // quants for 8 q8_0 blocks +} block_q8_0x8; +static_assert(sizeof(block_q8_0x8) == 8 * sizeof(ggml_half) + QK8_0 * 8, "wrong q8_0x8 block size/padding"); + +#define QK8_2 32 +typedef struct { + uint16_t d; + uint16_t s; + int8_t qs[QK8_2]; // quants +} block_q8_2; +static_assert(sizeof(block_q8_2) == sizeof(ggml_half) + sizeof(int16_t) + QK8_2, "wrong q8_2 block size/padding"); + +typedef struct { + uint16_t d[8]; + int8_t qs[4*QK8_2]; +} block_q8_2_x4; +static_assert(sizeof(block_q8_2_x4) == 4*sizeof(block_q8_2), "wrong q8_2_x4 block size/padding"); + +typedef struct { + ggml_half d[8]; + uint8_t scales[QK_K/4]; // scales and mins, quantized with 4 bits + uint8_t qs[QK_K]; // quants +} block_q2_k_r4; +static_assert(sizeof(block_q2_k_r4) == 8*sizeof(ggml_half) + QK_K/4 + QK_K, "wrong q2_k_r4 block size/padding"); + +typedef struct { + ggml_half d[4]; // super-block scales + uint8_t scales_h[QK_K/16]; // scales quantized with 6 bits (high 2 bits) + uint8_t scales_l[QK_K/8]; // scales quantized with 6 bits (low 4 bits) + uint8_t qh[QK_K/2]; // quants - high bit + uint8_t qs[QK_K]; // quants - low 2 bits +} block_q3_k_r4; +static_assert(sizeof(block_q3_k_r4) == 4*sizeof(ggml_half) + QK_K/16 + QK_K/8 + QK_K/2 + QK_K, "wrong q3_k_r4 block size/padding"); + +typedef struct { + ggml_half d[8]; + uint8_t scales_h[QK_K/16];// scales and mins, quantized with 6 bits + uint8_t scales_l[QK_K/8]; // scales and mins, quantized with 6 bits + uint8_t qs[QK_K*2]; // 4--bit quants +} block_q4_k_r4; +static_assert(sizeof(block_q4_k_r4) == 8*sizeof(ggml_half) + QK_K/16 + QK_K/8 + QK_K*2, "wrong q4_k_r4 block size/padding"); + +typedef struct { + ggml_half d[8]; + uint8_t scales_h[QK_K/16];// scales and mins, quantized with 6 bits + uint8_t scales_l[QK_K/8]; // scales and mins, quantized with 6 bits + uint8_t qh[QK_K/2]; // quants, high bit + uint8_t qs[QK_K*2]; // quants, low 4 bits +} block_q5_k_r4; +static_assert(sizeof(block_q5_k_r4) == 8*sizeof(ggml_half) + QK_K/16 + QK_K/8 + QK_K/2 + QK_K*2, "wrong q5_k_r4 block size/padding"); + +typedef struct { + ggml_half d[4]; // super-block scale + int8_t scales[QK_K/4]; // scales, quantized with 8 bits + uint8_t qh[QK_K]; // quants, upper 2 bits + uint8_t ql[QK_K*2]; // quants, lower 4 bits +} block_q6_k_r4; +static_assert(sizeof(block_q6_k_r4) == 4*sizeof(ggml_half) + QK_K/4 + 3*QK_K, "wrong q6_k_r4 block size/padding"); + +typedef struct { + float d; // delta + int8_t qs[64]; // quants +} block_q8_K64; +static_assert(sizeof(block_q8_K64) == sizeof(float) + 64, "wrong q8_K64 block size/padding"); + +typedef struct { + float d; // delta + int16_t bsums[4]; // quant sums for blocks of 32 + int8_t qs[128]; // quants +} block_q8_K128; +static_assert(sizeof(block_q8_K128) == sizeof(float) + 4*sizeof(int16_t) + 128, "wrong q8_K128 block size/padding"); + +typedef struct { + ggml_half d[8]; // delta + int8_t qs[8*QK_K]; // quants, stored as unsigned ints +} block_q8_k_r8; +static_assert(sizeof(block_q8_k_r8) == 8*sizeof(ggml_half) + 8*QK_K, "wrong q8_k_r8 block size/padding"); + +typedef struct { + ggml_half d[16]; // delta + int8_t qs[16*QK_K]; // quants, stored as unsigned ints +} block_q8_k_r16; +static_assert(sizeof(block_q8_k_r16) == 16*sizeof(ggml_half) + 16*QK_K, "wrong q8_k_r16 block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t sas[QK_K/2]; + uint8_t qs[QK_K/2]; +} block_iq2_xxs_r4; +static_assert(sizeof(block_iq2_xxs_r4) == 4*sizeof(block_iq2_xxs), "wrong iq2_xxs_r4 block size/padding"); + +typedef struct { + ggml_half d[4]; + uint16_t qs[QK_K/2]; + uint8_t scales[QK_K/8]; +} block_iq2_xs_r4; +static_assert(sizeof(block_iq2_xs_r4) == 4*sizeof(block_iq2_xs), "wrong iq2_xs_r4 block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t qs[QK_K/2]; + uint8_t qh[QK_K/8]; + uint8_t signs[QK_K/2]; + uint8_t scales[QK_K/8]; +} block_iq2_s_r4; +static_assert(sizeof(block_iq2_s_r4) == 4*sizeof(block_iq2_s), "wrong iq2_s_r4 block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t sas[QK_K/2]; + uint8_t qs[QK_K]; +} block_iq3_xxs_r4; +static_assert(sizeof(block_iq3_xxs_r4) == 4*sizeof(block_iq3_xxs), "wrong iq3_xxs_r4 block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t qs[QK_K]; + uint8_t qh[QK_K/8]; + uint8_t signs[QK_K/2]; + uint8_t scales[4*IQ3S_N_SCALE]; +} block_iq3_s_r4; +static_assert(sizeof(block_iq3_s_r4) == 4*sizeof(block_iq3_s), "wrong iq3_s_r4 block size/padding"); + +typedef struct { + uint8_t qs[16]; + uint16_t qh[4]; +} block_iq1_s_r4; +static_assert(sizeof(block_iq1_s_r4) == 24, "wrong iq1_s_r4 block size/padding"); + +// 1.75 bpw - blocks of 32 with 4 interleaved rows = 128 quants +typedef struct { + uint8_t qs[16]; // grid index, low 8 bits + uint8_t qh[ 8]; // grid index, high 3 bits + grid shift bits (for two groups of 8) + uint8_t scales[4]; // 4-bit block scales +} block_iq1_m_r4; +static_assert(sizeof(block_iq1_m_r4) == 28, "wrong iq1_m_r4 block size/padding"); + +// +// Bonsai +// +#define QK1_0_G128 128 +typedef struct { + ggml_half d; + uint8_t qs[QK1_0_G128 / 8]; +} block_q1_0_g128; +static_assert(sizeof(block_q1_0_g128) == sizeof(ggml_half) + QK1_0_G128 / 8, "wrong q1_0_g128 block size/padding"); + +// +// Bitnet and TriLM - implemented as 1.625 bpw +// +#define QK_IQ1BN 64 +typedef struct { + uint8_t ql[12]; + uint8_t extra; +} block_iq1_bn; +static_assert(sizeof(block_iq1_bn) == 13, "wrong iq1_bn block size/padding"); + +// +// Bitnet and TriLM - implemented as 2.0 bpw +// +#define QK_IQ2BN 64 +typedef struct { + uint8_t qs[QK_IQ2BN/4]; +} block_iq2_bn; +static_assert(sizeof(block_iq2_bn) == QK_IQ2BN/4, "wrong iq2_bn block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t qs[2*QK4_NL]; +} block_iq4_nl_r4; +static_assert(sizeof(block_iq4_nl_r4) == 4*sizeof(ggml_half) + 2*QK4_NL, "wrong iq4_nl_r4 block size/padding"); + +typedef struct { + ggml_half d[8]; + uint8_t qs[4*QK4_NL]; +} block_iq4_nl_r8; +static_assert(sizeof(block_iq4_nl_r8) == 8*sizeof(ggml_half) + 4*QK4_NL, "wrong iq4_nl_r8 block size/padding"); + +typedef struct { + ggml_half d[8]; + uint8_t scales_h[QK_K/16]; + uint8_t scales_l[QK_K/ 8]; + uint8_t qs[QK_K*4]; +} block_iq4_xs_r8; +static_assert(sizeof(block_iq4_xs_r8) == 8*sizeof(block_iq4_xs), "wrong iq4_xs_rs block size/padding"); + +typedef struct { + uint8_t scales[QK_K/32]; + uint8_t qs[QK_K/2]; +} block_iq4_ks; +static_assert(sizeof(block_iq4_ks) == QK_K/32 + QK_K/2, "wrong iq4_ks block size/padding"); + +typedef struct { + uint8_t scales[QK_K/8]; + uint8_t qs[QK_K*2]; +} block_iq4_ks_r4; +static_assert(sizeof(block_iq4_ks_r4) == 4*sizeof(block_iq4_ks), "wrong iq4_ks_r4 block size/padding"); + +typedef struct { + uint32_t qs[QK_K/8]; +} block_iq4_kss; +static_assert(sizeof(block_iq4_kss) == QK_K/8*sizeof(uint32_t), "wrong iq4_kss block size/padding"); + +typedef struct { + ggml_half d; + uint16_t extra; + uint8_t scales[QK_K/32]; + uint8_t qs[QK_K/4]; +} block_iq2_k; +static_assert(sizeof(block_iq2_k) == sizeof(ggml_half) + sizeof(uint16_t) + QK_K/32 + QK_K/4, "wrong iq2_k block size/padding"); + +typedef struct { + uint16_t scales_h; + uint8_t scales_l[QK_K/64]; + uint8_t qs[QK_K/4]; + uint8_t qh[QK_K/16]; +} block_iq2_kl; +static_assert(sizeof(block_iq2_kl) == sizeof(uint16_t) + QK_K/64 + QK_K/4 + QK_K/16, "wrong iq2_kl block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t extra[8]; + uint8_t scales[QK_K/8]; + uint8_t qs[QK_K]; +} block_iq2_k_r4; +static_assert(sizeof(block_iq2_k_r4) == 4*sizeof(block_iq2_k), "wrong iq2_k_r4 block size/padding"); + +typedef struct { + uint16_t extra; + uint8_t scales[QK_K/64]; + uint8_t qs[QK_K/4]; +} block_iq2_ks; +static_assert(sizeof(block_iq2_ks) == sizeof(uint16_t) + QK_K/64 + QK_K/4, "wrong iq2_ks block size/padding"); + +typedef struct { + uint8_t sh[QK_K/32]; // 4-bit scales + 13th bits for groups of 8 + uint8_t ql[QK_K/8]; // low 8 bits for groups of 8 + uint8_t qh[QK_K/16]; // high 4 bits for groups of 8 +} block_iq1_kt; +static_assert(sizeof(block_iq1_kt) == QK_K/8 + QK_K/16 + QK_K/32, "wrong iq1_kt block size/padding"); + +typedef struct { + uint8_t scales[QK_K/64]; + uint8_t ql[QK_K/4]; +} block_iq2_kt; +static_assert(sizeof(block_iq2_kt) == QK_K/4 + QK_K/64, "wrong iq2_kt block size/padding"); + +typedef struct { + uint8_t scales[QK_K/64]; + uint8_t ql[QK_K/4]; + uint8_t qh[QK_K/8]; +} block_iq3_kt; +static_assert(sizeof(block_iq3_kt) == QK_K/4 + QK_K/8 + QK_K/64, "wrong iq3_kt block size/padding"); + +typedef struct { + uint32_t qs[QK_K/8]; +} block_iq4_kt; +static_assert(sizeof(block_iq4_kt) == QK_K/2, "wrong iq4_kt block size/padding"); + +typedef struct { + ggml_half d; + uint16_t extra; + uint16_t scales_h; + uint8_t scales_l[QK_K/32]; + uint8_t qs[QK_K/4]; + uint8_t qh[QK_K/8]; +} block_iq3_k; +static_assert(sizeof(block_iq3_k) == sizeof(ggml_half) + 2*sizeof(uint16_t) + QK_K/32 + QK_K/4 + QK_K/8, "wrong iq3_k block size/padding"); + +typedef struct { + uint16_t extra; + uint8_t scales[QK_K/64]; + uint8_t qs[QK_K/4]; + uint8_t qh[QK_K/8]; +} block_iq3_ks; +static_assert(sizeof(block_iq3_ks) == sizeof(uint16_t) + QK_K/64 + QK_K/4 + QK_K/8, "wrong iq3_ks block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t extra[8]; + uint8_t scales_h[QK_K/32]; + uint8_t scales_l[QK_K/8]; + uint8_t qs[QK_K]; + uint8_t qh[QK_K/2]; +} block_iq3_k_r4; +static_assert(sizeof(block_iq3_k_r4) == 4*sizeof(block_iq3_k), "wrong iq3_k_r4 block size/padding"); + +typedef struct { + ggml_half d; + uint16_t extra; + uint8_t scales_h[QK_K/64]; + uint8_t scales_l[QK_K/32]; + uint8_t qs[QK_K/2]; +} block_iq4_k; +static_assert(sizeof(block_iq4_k) == sizeof(ggml_half) + sizeof(uint16_t) + QK_K/2 + 3*QK_K/64, "wrong iq4_k block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t extra[8]; + uint8_t scales_h[QK_K/16]; + uint8_t scales_l[QK_K/8]; + uint8_t qs[QK_K*2]; +} block_iq4_k_r4; +static_assert(sizeof(block_iq4_k_r4) == 4*sizeof(block_iq4_k), "wrong iq4_k_r4 block size/padding"); + +typedef struct { + ggml_half d; + uint16_t extra; + uint8_t scales_h[QK_K/64]; + uint8_t scales_l[QK_K/32]; + uint8_t qs[QK_K/2]; + uint8_t qh[QK_K/8]; +} block_iq5_k; +static_assert(sizeof(block_iq5_k) == sizeof(ggml_half) + sizeof(uint16_t) + QK_K/2 + QK_K/8 + 3*QK_K/64, "wrong iq5_k block size/padding"); + +typedef struct { + ggml_half d[4]; + uint8_t extra[8]; + uint8_t scales_h[QK_K/16]; + uint8_t scales_l[QK_K/8 ]; + uint8_t qs[QK_K*2]; + uint8_t qh[QK_K/2]; +} block_iq5_k_r4; +static_assert(sizeof(block_iq5_k_r4) == 4*sizeof(block_iq5_k), "wrong iq5_k_r4 block size/padding"); + +typedef struct { + ggml_half d; + uint16_t extra; + int8_t scales[QK_K/16]; + uint8_t qs[QK_K/2]; + uint8_t qh[QK_K/4]; +} block_iq6_k; +static_assert(sizeof(block_iq6_k) == sizeof(ggml_half) + sizeof(uint16_t) + QK_K/2 + QK_K/4 + QK_K/16, "wrong iq6_k block size/padding"); + +typedef struct { + uint8_t scales[QK_K/32]; + uint8_t qs[QK_K/2]; + uint8_t qh[QK_K/8]; +} block_iq5_ks; +static_assert(sizeof(block_iq5_ks) == QK_K/32 + QK_K/2 + QK_K/8, "wrong iq5_ks block size/padding"); + +typedef struct { + uint8_t scales[QK_K/8]; + uint8_t qs[QK_K*2]; + uint8_t qh[QK_K/2]; +} block_iq5_ks_r4; +static_assert(sizeof(block_iq5_ks_r4) == 4*sizeof(block_iq5_ks), "wrong iq5_ks_r4 block size/padding"); + + +// opencoti F5-opt W2 (#290): ik_llama's block_q8_K carries an extra `float sum;` +// after `d` (used only for intermediate quantization / dot products, never GGUF- +// serialized). llamafile's GGUF-frozen block_q8_K has no such member and must not +// be edited. iqk_quantize_row_q8_K_T() writes y[i].sum, so we give it an IK-private +// struct with the ik_llama layout to cast to. Body verbatim from ik_llama +// ggml-common.h; static_assert pins the (larger) ik layout. +typedef struct { + float d; // delta + float sum; // sum of quants in the entire block + int8_t qs[QK_K]; // quants + int16_t bsums[QK_K/16]; // sum of quants in groups of 16 +} iqk_block_q8_K; +static_assert(sizeof(iqk_block_q8_K) == 2*sizeof(float) + QK_K + QK_K/16*sizeof(int16_t), "wrong iqk_block_q8_K size/padding"); + +// IK-only lookup table (ggml-common.h IMPL delta); static = per-TU copy, harmless. +static const int8_t iq4k_values[32] = { +-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113, + -123, -100, -79, -61, -45, -31, -18, -6, 5, 17, 29, 42, 57, 73, 93, 117 +}; + +// opencoti F5-opt W2 (#290): more IK-only lookup tables (ggml-common.h IMPL deltas, +// needed by iqk_quantize.cpp). static = per-TU copy, harmless. +static const int8_t iq2nl_values[8] = { + -31, -13, 1, 17, -26, -8, 6, 22 +}; + +static const uint16_t iq2kl_values[32] = { + 0xe9c1, 0x0dc1, 0xc1d8, 0xf6d8, 0x0dd8, 0x2fd8, 0xd8e9, 0xe9e9, 0x01e9, 0x0de9, 0x1ce9, 0xc1f6, 0x01f6, 0x0df6, 0x2ff6, 0xe901, + 0xf601, 0x0101, 0x0d01, 0x1c01, 0xd80d, 0xe90d, 0xf60d, 0x010d, 0x0d0d, 0xc11c, 0xe91c, 0x011c, 0x1c1c, 0x2f1c, 0xe92f, 0x0d2f, +}; + +static const int8_t iq3nl_values[16] = { + -63, -40, -23, -10, 1, 13, 28, 47, + -59, -36, -19, -6, 5, 17, 32, 51, +}; + +static const int8_t iq5nl_values[64] = { + -126, -114, -103, -92, -83, -74, -65, -57, -50, -43, -36, -30, -24, -18, -12, -6, -1, 5, 11, 17, 23, 29, 36, 43, 51, 59, 68, 77, 87, 97, 109, 121, + -124, -112, -101, -90, -81, -72, -63, -55, -48, -41, -34, -28, -22, -16, -10, -4, 1, 7, 13, 19, 25, 31, 38, 45, 53, 61, 70, 79, 89, 99, 111, 123, +}; + +static const int8_t iq6nl_values[128] = { + -127, -121, -115, -109, -104, -98, -93, -88, -84, -79, -74, -70, -66, -62, -58, -54, + -51, -47, -44, -40, -37, -34, -31, -28, -25, -22, -19, -16, -13, -11, -8, -5, + -2, 0, 3, 6, 9, 12, 14, 17, 20, 23, 27, 30, 33, 36, 40, 44, + 47, 51, 55, 59, 63, 68, 72, 77, 82, 87, 92, 98, 103, 109, 115, 121, + -126, -120, -114, -108, -103, -97, -92, -87, -83, -78, -73, -69, -65, -61, -57, -53, + -50, -46, -43, -39, -36, -33, -30, -27, -24, -21, -18, -15, -12, -10, -7, -4, + -1, 1, 4, 7, 10, 13, 15, 18, 21, 24, 28, 31, 34, 37, 41, 45, + 48, 52, 56, 60, 64, 69, 73, 78, 83, 88, 93, 99, 104, 110, 116, 122, +}; + +#endif // IQK_COMMON_EXTRA_H diff --git a/llama.cpp/ggml/src/iqk/iqk_config.h b/llama.cpp/ggml/src/iqk/iqk_config.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_config.h @@ -0,0 +1,101 @@ +// +// Copyright (C) 2024-2025 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// + +#pragma once + +#if defined IQK_IMPLEMENT +#undef IQK_IMPLEMENT +#endif + +#if defined __AVX2__ || defined __ARM_FEATURE_DOTPROD +#define IQK_IMPLEMENT +#endif + +// opencoti F5-opt W2 (#290): ik_llama's iqk headers (iqk_common.h / iqk_utils.h) +// use __m256/__m128i in early inline code that, upstream, relies on ik_llama's +// ggml-impl.h having pulled the SIMD intrinsics header first. llamafile's +// ggml-impl.h does NOT (it surfaces intrinsics via ggml-cpu/ggml-cpu-impl.h, +// which the iqk TUs don't include), so we pull them here — the root iqk config +// header every iqk TU includes first. Guarded to the active ISA; non-x86/non-ARM +// or non-AVX2 builds get nothing (IQK_IMPLEMENT is also off there). +#if defined IQK_IMPLEMENT +#if defined __x86_64__ || defined __i386__ +#include +#elif defined __aarch64__ || defined __ARM_NEON +#include +#endif +#endif + +// opencoti F5-opt W2 (#290): ggml_vdotq_s32 — ik_llama defines it in ggml-impl.h +// (which the iqk TUs include), but llamafile MOVED it to ggml-cpu/ggml-cpu-impl.h +// (which the iqk TUs do NOT include). iqk_quantize.cpp's #ifdef __ARM_NEON path +// calls it, so under cosmocc's aarch64 slice (NEON on, DOTPROD off) the symbol is +// missing. Re-supply the upstream-identical NEON fallback here, mirroring the +// intrinsics shim above. Guarded by __ARM_NEON + sentinel so we never collide with +// llamafile's ggml-cpu-impl.h should a TU ever pull it. +#if defined __ARM_NEON && !defined GGML_VDOTQ_S32_DEFINED +#define GGML_VDOTQ_S32_DEFINED +#include +#if !defined(__ARM_FEATURE_DOTPROD) +inline static int32x4_t ggml_vdotq_s32(int32x4_t acc, int8x16_t a, int8x16_t b) { + const int16x8_t p0 = vmull_s8(vget_low_s8 (a), vget_low_s8 (b)); + const int16x8_t p1 = vmull_s8(vget_high_s8(a), vget_high_s8(b)); + return vaddq_s32(acc, vaddq_s32(vpaddlq_s16(p0), vpaddlq_s16(p1))); +} +#else +#define ggml_vdotq_s32(a, b, c) vdotq_s32(a, b, c) +#endif // !defined(__ARM_FEATURE_DOTPROD) +#endif // __ARM_NEON && !GGML_VDOTQ_S32_DEFINED + +#ifdef GGML_SHARED +# if defined(_WIN32) && !defined(__MINGW32__) +# ifdef GGML_BUILD +# define IQK_API __declspec(dllexport) +# else +# define IQK_API __declspec(dllimport) +# endif +# else +# define IQK_API __attribute__ ((visibility ("default"))) +# endif +#else +# define IQK_API +#endif + +#ifdef _MSC_VER +#define IQK_NOINLINE __declspec(noinline) +#define IQK_ALWAYS_INLINE inline +#if !defined __x86_64__ && defined _M_X64 +#define __x86_64__ +#endif +#else +#define IQK_NOINLINE __attribute__((__noinline__)) +#define IQK_ALWAYS_INLINE __attribute__((__always_inline__)) +#endif + +#if defined __x86_64__ +#if defined HAVE_FANCY_SIMD + #undef HAVE_FANCY_SIMD +#endif +#if defined(__AVX512F__) && defined(__AVX512VNNI__) && defined(__AVX512VL__) && defined(__AVX512BW__) && defined(__AVX512DQ__) + #define HAVE_FANCY_SIMD +#endif +#if defined HAVE_VNNI256 + #undef HAVE_VNNI256 +#endif +#if defined(__AVXVNNI__) || (defined(__AVX512VNNI__) && defined(__AVX512VL__)) + #define HAVE_VNNI256 +#endif +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + #define ggml_mm256_dpbusd_epi32 _mm256_dpbusd_epi32 + #define ggml_mm256_dpwssd_epi32 _mm256_dpwssd_epi32 + #define ggml_mm_dpbusd_epi32 _mm_dpbusd_epi32 +#elif defined(__AVXVNNI__) + #define ggml_mm256_dpbusd_epi32 _mm256_dpbusd_avx_epi32 + #define ggml_mm256_dpwssd_epi32 _mm256_dpwssd_avx_epi32 + #define ggml_mm_dpbusd_epi32 _mm_dpbusd_avx_epi32 +#endif +#endif + diff --git a/llama.cpp/ggml/src/iqk/iqk_fa_dispatch.cpp b/llama.cpp/ggml/src/iqk/iqk_fa_dispatch.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_fa_dispatch.cpp @@ -0,0 +1,230 @@ +// -*- mode:c++;indent-tabs-mode:nil;c-basic-offset:4;coding:utf-8 -*- +// vi: set et ft=cpp fenc=utf-8 :vi +// +// opencoti F5-opt W2 (#290) — CPU Flash-Attention dispatch + K-repack. +// +// VENDORED (extracted) from ik_llama.cpp +// repo: https://github.com/ikawrakow/ik_llama.cpp +// commit: 8960c5ba5ee9db30ba838304373aa4dbec9f7cbd +// file: ggml/src/iqk/iqk_mul_mat.cpp lines 1249-1433 +// (the `#ifdef GGML_IQK_FLASH_ATTENTION` block: iqk_repack_k + +// iqk_flash_attn_impl, plus its inner template includes) +// license: MIT — Copyright (C) 2024 Iwan Kawrakow +// +// WHY THIS FILE EXISTS (do not "fix" by merging back): +// The two functions below are the ONLY parts of the 1780-LOC iqk_mul_mat.cpp +// the FA engine needs. The rest of that TU is the universal GEMM dispatcher, +// which (a) #includes + references the iquants/1bit/iqk_quants/ktquants GEMM +// families (~800 KB of extra .cpp we do NOT want), and (b) defines the +// extern-C symbols `iqk_mul_mat` / `iqk_mul_mat_moe` that COLLIDE with +// llamafile's own tinyBLAS `iqk_mul_mat.inc`. Extracting just these two +// functions lets us drop iqk_mul_mat.cpp entirely — no GEMM-family balloon, +// no ODR collision, no symbol rename needed. +// +// To re-sync: re-extract lines 1249-1433 of upstream iqk_mul_mat.cpp at the +// pinned SHA above (boundaries = the GGML_IQK_FLASH_ATTENTION guard pair). +// +#include "iqk_config.h" + +#if defined IQK_IMPLEMENT + +#include +#include +#include +#include + +#include "ggml-impl.h" +#include "ggml-quants.h" +#include "ggml-common.h" +#include "iqk_mul_mat.h" +#include "iqk_quantize.h" +#include "iqk_utils.h" +#include "iqk_common.h" + +#ifdef GGML_IQK_FLASH_ATTENTION + +void * iqk_repack_k(int int_type_k, int nek0, int nek1, int nek2, int nek3, long nbk1, long nbk2, long nbk3, + const void * data, void * work, int ith, int nth, int& repacked_type, uint64_t& row_size) { + repacked_type = int_type_k; + auto type_k = ggml_type(int_type_k); + if (type_k != GGML_TYPE_Q8_0 || nek0%QK8_0 != 0) return work; + int nrows = nek1*nek2*nek3; + if (nrows%8 != 0) return work; + repacked_type = int(GGML_TYPE_Q8_0_R8); + row_size = ggml_row_size(GGML_TYPE_Q8_0, nek0); + void * result = (char *)work + nrows*row_size; + int npt = 8*((nrows/8 + nth - 1)/nth); + int first = npt*ith; + if (first >= nrows) return result; + int last = std::min(first + npt, nrows); + const block_q8_0 * x8[8]; + auto y = (block_q8_0_r8 *)((char *)work + first*row_size); + int nblock = nek0/QK8_0; +#ifdef __ARM_NEON + int8x16x2_t m0, m1, m2, m3; +#endif + for (int row = first; row < last; row += 8) { + int ik3 = row/(nek1*nek2); + int ik2 = (row - ik3*nek1*nek2)/nek1; + int ik1 = row - ik3*nek1*nek2 - ik2*nek1; + auto this_data = (const char *)data + ik1*nbk1 + ik2*nbk2 + ik3*nbk3; + for (int k = 0; k < 8; ++k) x8[k] = (const block_q8_0 *)(this_data + k*nbk1); + for (int ib = 0; ib < nblock; ++ib) { + for (int k = 0; k < 8; ++k) y[ib].d[k] = x8[k][ib].d; +#ifdef __AVX2__ + auto m0 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[4][ib].qs), _mm_loadu_si128((const __m128i *)x8[0][ib].qs)); + auto m1 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[5][ib].qs), _mm_loadu_si128((const __m128i *)x8[1][ib].qs)); + auto m2 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[6][ib].qs), _mm_loadu_si128((const __m128i *)x8[2][ib].qs)); + auto m3 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[7][ib].qs), _mm_loadu_si128((const __m128i *)x8[3][ib].qs)); + auto t0 = _mm256_unpacklo_epi32(m0, m1); + auto t1 = _mm256_unpacklo_epi32(m2, m3); + auto t2 = _mm256_unpackhi_epi32(m0, m1); + auto t3 = _mm256_unpackhi_epi32(m2, m3); + m0 = _mm256_unpacklo_epi64(t0, t1); + m1 = _mm256_unpackhi_epi64(t0, t1); + m2 = _mm256_unpacklo_epi64(t2, t3); + m3 = _mm256_unpackhi_epi64(t2, t3); + //#ifdef HAVE_FANCY_SIMD + // m0 = _mm256_add_epi8(m0, _mm256_set1_epi8(127)); + // m1 = _mm256_add_epi8(m1, _mm256_set1_epi8(127)); + // m2 = _mm256_add_epi8(m2, _mm256_set1_epi8(127)); + // m3 = _mm256_add_epi8(m3, _mm256_set1_epi8(127)); + //#endif + _mm256_storeu_si256((__m256i *)y[ib].qs + 0, m0); + _mm256_storeu_si256((__m256i *)y[ib].qs + 1, m1); + _mm256_storeu_si256((__m256i *)y[ib].qs + 2, m2); + _mm256_storeu_si256((__m256i *)y[ib].qs + 3, m3); + m0 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[4][ib].qs+1), _mm_loadu_si128((const __m128i *)x8[0][ib].qs+1)); + m1 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[5][ib].qs+1), _mm_loadu_si128((const __m128i *)x8[1][ib].qs+1)); + m2 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[6][ib].qs+1), _mm_loadu_si128((const __m128i *)x8[2][ib].qs+1)); + m3 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[7][ib].qs+1), _mm_loadu_si128((const __m128i *)x8[3][ib].qs+1)); + t0 = _mm256_unpacklo_epi32(m0, m1); + t1 = _mm256_unpacklo_epi32(m2, m3); + t2 = _mm256_unpackhi_epi32(m0, m1); + t3 = _mm256_unpackhi_epi32(m2, m3); + m0 = _mm256_unpacklo_epi64(t0, t1); + m1 = _mm256_unpackhi_epi64(t0, t1); + m2 = _mm256_unpacklo_epi64(t2, t3); + m3 = _mm256_unpackhi_epi64(t2, t3); + //#ifdef HAVE_FANCY_SIMD + // m0 = _mm256_add_epi8(m0, _mm256_set1_epi8(127)); + // m1 = _mm256_add_epi8(m1, _mm256_set1_epi8(127)); + // m2 = _mm256_add_epi8(m2, _mm256_set1_epi8(127)); + // m3 = _mm256_add_epi8(m3, _mm256_set1_epi8(127)); + //#endif + _mm256_storeu_si256((__m256i *)y[ib].qs + 4, m0); + _mm256_storeu_si256((__m256i *)y[ib].qs + 5, m1); + _mm256_storeu_si256((__m256i *)y[ib].qs + 6, m2); + _mm256_storeu_si256((__m256i *)y[ib].qs + 7, m3); +#elif defined __ARM_NEON + for (int l = 0; l < 2; ++l) { + m0.val[0] = vld1q_s8(x8[0][ib].qs+16*l); m0.val[1] = vld1q_s8(x8[4][ib].qs+16*l); + m1.val[0] = vld1q_s8(x8[1][ib].qs+16*l); m1.val[1] = vld1q_s8(x8[5][ib].qs+16*l); + m2.val[0] = vld1q_s8(x8[2][ib].qs+16*l); m2.val[1] = vld1q_s8(x8[6][ib].qs+16*l); + m3.val[0] = vld1q_s8(x8[3][ib].qs+16*l); m3.val[1] = vld1q_s8(x8[7][ib].qs+16*l); + auto row01 = vtrnq_s32(vreinterpretq_s32_s8(m0.val[0]), vreinterpretq_s32_s8(m1.val[0])); + auto row23 = vtrnq_s32(vreinterpretq_s32_s8(m2.val[0]), vreinterpretq_s32_s8(m3.val[0])); + m0.val[0] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m1.val[0] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + m2.val[0] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m3.val[0] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + row01 = vtrnq_s32(vreinterpretq_s32_s8(m0.val[1]), vreinterpretq_s32_s8(m1.val[1])); + row23 = vtrnq_s32(vreinterpretq_s32_s8(m2.val[1]), vreinterpretq_s32_s8(m3.val[1])); + m0.val[1] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m1.val[1] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + m2.val[1] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m3.val[1] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + vst1q_s8_x2(y[ib].qs + 0 + 128*l, m0); + vst1q_s8_x2(y[ib].qs + 32 + 128*l, m1); + vst1q_s8_x2(y[ib].qs + 64 + 128*l, m2); + vst1q_s8_x2(y[ib].qs + 96 + 128*l, m3); + } +#else + for (int l = 0; l < 4; ++l) { + for (int k = 0; k < 8; ++k) for (int i = 0; i < 4; ++i) { + y[ib].qs[32*l+4*k+i+ 0] = x8[k][ib].qs[i+4*l+ 0]; + y[ib].qs[32*l+4*k+i+128] = x8[k][ib].qs[i+4*l+16]; + } + } +#endif + } + y += nblock; + } + return result; +} + +#include "iqk_flash_impl.h" +#include "fa/iqk_fa_templates.h" + +bool iqk_flash_attn_impl(int int_type_k, // type of k + int int_type_v, // type of v + int Dk, // K head size + int Dv, // V head size + int nq1, // number of columns in q + int nk1, // number of rows in k + int stride_q, // distance between q columns in bytes + int stride_k, // distance between k rows in bytes + int stride_v, // distance between v rows in bytes + int stride_m, // distance between mask rows (in bytes + int stride_qkv, // distance between rows in mask (in bytes) + const float * q, // q matrix. + const void * k, // k matrix. Assumed to be fp16, nq x nk elements + const void * v, // v matrix. Assumed to be fp16, nq x nk elements + const void * mask, // mask. If not null, assumed to be fp16. nq x nk elements + const float * sinksf, // mask. If not null, assumed to be fp16. nq x nk elements + [[maybe_unused]] int nsinks, + float scale, // scale applied before softmax + float softcap, // if > 0, a "soft-cap" operation is applied before softmax + float * qkv, // v*softmax(scale*(k*q)) + float * M, float * S) { + + if (!mask || nk1%32 != 0) return false; // the implementation assumes mask is not null and nk is a multiple of 32 + + if (Dk == 576 && Dv == 512) { + return iqk_fa_576_512(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + if (Dk == 512 && Dv == 512) { + return iqk_fa_512_512(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + if (Dk == 320 && Dv == 256) { + return iqk_fa_320_256(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + + if (Dk == 192 && Dv == 128) { + return iqk_fa_192_128(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + + if (Dk == 192 && Dv == 192) { + return iqk_fa_192_192(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + + if (Dk == 256 && Dv == 256) { + return iqk_fa_256_256(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + + if (Dk == 128 && Dv == 128) { + return iqk_fa_128_128(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + + if (Dk == 96 && Dv == 96) { + return iqk_fa_96_96(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + + if (Dk == 64 && Dv == 64) { + return iqk_fa_64_64(int_type_k, int_type_v, nq1, nk1, stride_q, stride_k, stride_v, stride_m, stride_qkv, + q, k, v, mask, scale, softcap, qkv, sinksf, M, S); + } + + return false; +} +#endif + +#endif // IQK_IMPLEMENT diff --git a/llama.cpp/ggml/src/iqk/iqk_flash_attn.cpp b/llama.cpp/ggml/src/iqk/iqk_flash_attn.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_flash_attn.cpp @@ -0,0 +1,525 @@ +// +// Copyright (C) 2024-2025 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// + +#include "iqk_config.h" +#include "iqk_mul_mat.h" +#include "iqk_flash_impl.h" +#include "ggml.h" +#include "iqk_ggml_type_ext.h" // opencoti F5-opt W2 (#290): ik_llama ggml_type enum delta (this TU includes ggml.h directly, not the ggml-common.h shim) + +#if defined IQK_IMPLEMENT && defined GGML_IQK_FLASH_ATTENTION + +#include +#include +#include +#include +#include +#include +#include + +namespace { +inline uint32_t simple_gcd(uint32_t a, uint32_t b) { + while (a != b) { + if (a > b) a -= b; + else b -= a; + } + return a; +} +inline void accumulate_qkv(int Dv, float& M, float& S, float Mj, float Sj, float * Racc, const float * R) { + if (Mj == -INFINITY) return; + if (Mj > M) { + if (M == -INFINITY) { + std::memcpy(Racc, R, Dv*sizeof(float)); + S = Sj; + } else { + float c = exp(M - Mj); + S = c*S + Sj; + for (int i = 0; i < Dv; ++i) Racc[i] = c*Racc[i] + R[i]; + } + M = Mj; + } else { + float c = exp(Mj - M); + S += c*Sj; + for (int i = 0; i < Dv; ++i) Racc[i] += c*R[i]; + } +} +} + +size_t iqk_fa_work_buffer_size(const struct ggml_tensor * dst, int nth) { + auto Q = dst->src[0]; + auto K = dst->src[1]; + auto V = dst->src[2]; + int rk2 = Q->ne[2]/K->ne[2]; + size_t size = 0; + if (Q->ne[1] >= 8 && K->type == GGML_TYPE_Q8_0) { + size = ggml_row_size(GGML_TYPE_Q8_0, K->ne[0]) * K->ne[1]*K->ne[2]*K->ne[3]; + } + if (Q->ne[1] == 1 && Q->ne[3] == 1 && Q->ne[2]/K->ne[2] > 1 && nth >= 1 && K->ne[1]/32 > 1) { + if (K->ne[2] > 1) { + int gcd = simple_gcd(K->ne[2], nth); + int nth_k = nth/gcd; + int nek2_k = K->ne[2]/gcd; + int nchunk = nek2_k*K->ne[1]/32; + int npt = (nchunk + nth_k - 1)/nth_k; + int nk; + if (npt*nth_k == nchunk) { + nk = 32 * (K->ne[1]*K->ne[2]/(32*nth)); + } else { + //int nm = std::max(1, npt/8); + int nm = 1; + while (true) { + if (nm*4 >= npt) break; + nm *= 2; + } + nk = 32*nm; + } + int nkk = (K->ne[1] + nk - 1)/nk; + int nstep_k = K->ne[2]*nkk; + size_t result_size = (V->ne[0] + 16)*Q->ne[2]/K->ne[2]*sizeof(float); + size += nstep_k*result_size; + return size; + } + int nstep_k = K->ne[1]/32; + if (nstep_k >= 4*nth) { + auto size_thread = (V->ne[0] + 16)*rk2*sizeof(float); + size += size_thread*nth; + return size; + } + int gcd_k = simple_gcd(nstep_k, nth); + if (gcd_k >= 1) { + int nth_k = nth/gcd_k; + int nq_per_thread = (rk2 + nth_k - 1)/nth_k; + if (nq_per_thread > 1) { + auto size_thread = (V->ne[0] + 16)*nq_per_thread*sizeof(float); + size += size_thread*nth; + return size; + } + } + int rv2 = Q->ne[2] / V->ne[2]; + if (Q->ne[1] == 1 && Q->ne[3] == 1 && rk2 > 1 && rk2 == rv2 && K->ne[1]*K->ne[2] >= 32*nth) { + auto result_size = (V->ne[0] + 16)*rk2*sizeof(float); + size += result_size*nth; + } + return size; + } + return size; +} + +static inline const std::unordered_set & supported_kv_types() { +#ifdef GGML_IQK_FA_ALL_QUANTS + static std::unordered_set k_supported = { + GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q8_KV, GGML_TYPE_Q6_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_IQ4_NL + }; +#else + static std::unordered_set k_supported = { + GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q8_KV, GGML_TYPE_Q6_0, + }; +#endif + return k_supported; +} + +static inline bool are_kv_types_supported(ggml_type type_k, ggml_type type_v) { + if (type_k == GGML_TYPE_BF16) { + if (type_v != type_k) { + return false; + } +#ifdef __AVX512BF16__ + return true; +#else + return false; +#endif + } + auto & supported = supported_kv_types(); + auto it_k = supported.find(type_k); + auto it_v = supported.find(type_v); + return it_k != supported.end() && it_v != supported.end(); +} + +// TODO: get the ggml_type enum here without polution +// +extern "C" IQK_API bool iqk_flash_attn_noalibi(int type_q, int type_mask, float max_bias, + int neq3, int neq2, long nbq3, long nbq2, + int nek3, int nek2, long nbk3, long nbk2, + int nev3, int nev2, long nbv3, long nbv2, + int ne2, int ne1, long nb1, + int int_type_k_in, // type of k + int int_type_v, // type of v + int Dk, // K head size + int Dv, // V head size + int neq1, // number of columns in q + int nek1, // number of rows in k + int stride_q, // distance between q columns in bytes + int stride_k, // distance between k rows in bytes + int stride_v, // distance between v rows in bytes + int stride_m, // distance between mask rows (in bytes + const void * q, // q matrix. + const void * k, // k matrix. Assumed to be fp16, nq x nk elements + const void * v, // v matrix. Assumed to be fp16, nq x nk elements + const void * mask, // mask. If not null, assumed to be fp16. nq x nk elements + const void * sinks, // mask. If not null, assumed to be fp16. nq x nk elements + float scale, // scale applied before softmax + float softcap, // if > 0, a "soft-cap" operation is applied before softmax + float * qkv, // v*softmax(scale*(k*q)) + [[maybe_unused]] void * work_buffer_in, [[maybe_unused]] barrier_t barrier, [[maybe_unused]] void * barrier_data, + int ith, int nth, int n_swa) { + + if (type_q != 0 || type_mask != 1 || max_bias > 0) return false; + + if (auto type_k = ggml_type(int_type_k_in), type_v = ggml_type(int_type_v); !are_kv_types_supported(type_k, type_v)) { + if (ith == 0) { + fprintf(stderr, "\n==================== K cache %s coupled with V cache %s is not a supported combination on the CPU backend.\n", + ggml_type_name(type_k), ggml_type_name(type_v)); + auto & supported = supported_kv_types(); + fprintf(stderr, "Supported types are:\n"); + for (auto type : supported) { + fprintf(stderr, " %s\n", ggml_type_name(type)); + } + fprintf(stderr, " Warning: ik_llama.cpp does not support Q5_0 or Q5_1 KV cache on the CPU.\n"); +#ifdef __AVX512BF16__ + fprintf(stderr, " %s, but only if K and V are both %s\n", ggml_type_name(GGML_TYPE_BF16), ggml_type_name(GGML_TYPE_BF16)); +#endif +#ifndef GGML_IQK_FA_ALL_QUANTS + fprintf(stderr, " To enable q4_0, q4_1, and iq4_nl KV cache types, recompile with -DGGML_IQK_FA_ALL_QUANTS=ON\n"); +#endif + } + barrier(barrier_data); + GGML_ABORT("Fatal error"); + } + + if (n_swa > 0 && mask) { + constexpr int kMinBatch = 256; + int ntokens = std::max(kMinBatch, neq1); + int nblock = (ntokens + n_swa + kMinBatch - 1)/kMinBatch; + int first = nek1 - nblock*kMinBatch; + if (first > 0) { + k = (const char *)k + int64_t(first)*stride_k; + v = (const char *)v + int64_t(first)*stride_v; + mask = (const uint16_t *)mask + first; + nek1 -= first; + } + } + + int rk2 = neq2/nek2; + int rv2 = neq2/nev2; + int rk3 = neq3/nek3; + int rv3 = neq3/nev3; + + int first_k = 0, last_k = nek1; + if (neq3 == 1 && rk2 > 1 && neq1 == 1 && nek1 > 256 && mask) { + // This is a quick hack for SWA models. + // Given that the mask is the same for all layers, ideally we should determine the + // cache bounds once, and reuse for the whole graph. But even with this simple hack + // we get non-negligible performance gains for SWA models and long context. + auto umask = (const uint16_t *)mask; + for (; first_k < last_k; ++first_k) { + if (umask[first_k] == 0) break; + } + for (; last_k > first_k; --last_k) { + if (umask[last_k-1] == 0) break; + } + int non = 32*((last_k - first_k + 31)/32); + first_k = std::max(0, last_k - non); + last_k = std::min(first_k + non, nek1); + //printf("nek1 = %d, first = %d, last = %d\n", nek1, first, last); + if (last_k - first_k <= 3*nek1/4 && (last_k - first_k)%32 == 0) { + //printf("Reducing from %d to %d\n", nek1, last_k - first_k); + k = (const void *)((const char *)k + first_k*stride_k); + v = (const void *)((const char *)v + first_k*stride_v); + mask = (const void *)((const uint16_t *)mask + first_k); + nek1 = last_k - first_k; + } + } + + int int_type_k = int_type_k_in; + auto work_buffer = work_buffer_in; + if (neq1 >= 8) { + uint64_t row_size = 0; + work_buffer = iqk_repack_k(int_type_k, Dk, nek1, nek2, nek3, stride_k, nbk2, nbk3, k, work_buffer_in, ith, nth, int_type_k, row_size); + if (int_type_k != int_type_k_in) { + stride_k = row_size; + nbk2 = stride_k*nek1; + nbk3 = nbk2*nek2; + k = work_buffer_in; + barrier(barrier_data); + } + } + //uint64_t row_size = 0; + //auto work_buffer = iqk_repack_k(int_type_k, Dk, nek1, nek2, nek3, stride_k, nbk2, nbk3, k, work_buffer_in, ith, nth, int_type_k, row_size); + //if (int_type_k != int_type_k_in) { + // stride_k = row_size; + // nbk2 = stride_k*nek1; + // nbk3 = nbk2*nek2; + // k = work_buffer_in; + // barrier(barrier_data); + //} + + // Getting confused all the time about where to load data from and store the results to + // (especially when combining the results from the threads). + // So, for now, making it work just for MLA (nek2 = 1). + // I think it would also speed up things for GQA, but I'm leaving this for another day. + if (neq3 == 1 && rk2 > 1 && neq1 == 1 && nth >= 1 && nek1/32 > 1 && nek2 == 1) { + int nstep_k = nek1/32; + if (nstep_k >= 4*nth) { + int nstep_k_per_thread = (nstep_k + nth - 1)/nth; + int ith_mid = nth; + int nstep_k_this_thread = nstep_k_per_thread; + if (nstep_k_per_thread*nth > nstep_k) { + ith_mid = nstep_k - nth*(nstep_k_per_thread - 1); + if (ith >= ith_mid) --nstep_k_this_thread; + } + //if (ith == 0) fprintf(stderr, "nstep_k = %d, nstep_k_per_thread = %d, ith_mid = %d\n", nstep_k, nstep_k_per_thread, ith_mid); + nstep_k_per_thread *= 32; + nstep_k_this_thread *= 32; + + auto kv_offset = ith <= ith_mid ? ith*nstep_k_per_thread + : ith_mid*nstep_k_per_thread + (ith - ith_mid)*nstep_k_this_thread; + auto kth = (const char *)k + kv_offset*stride_k; + auto vth = (const char *)v + kv_offset*stride_v; + auto qth = (const char *)q; + auto mth = mask ? (const char *)mask + kv_offset*sizeof(uint16_t) : nullptr; // we don't have ggml_half available here + + auto work = (char *)work_buffer; + auto size_thread = (Dv + 16)*rk2*sizeof(float); + auto result_buffer = work; + auto work_this_thread = (float *)(result_buffer + ith*size_thread); + if (!iqk_flash_attn_impl(int_type_k, int_type_v, + Dk, Dv, rk2, nstep_k_this_thread, nbq2, stride_k, stride_v, 0, Dv, //Dk*sizeof(uint16_t), Dv, + (const float *)qth, (const void *)kth, (const void *)vth, (const void *)mth, nullptr, 0, + scale, softcap, + work_this_thread, work_this_thread + (Dv+0)*rk2, work_this_thread + (Dv+1)*rk2)) return false; + + barrier(barrier_data); + + for (int j = ith; j < rk2; j += nth) { + auto Racc = qkv + j*nb1/sizeof(float); + float M = -INFINITY, S = 0; + for (int jth = 0; jth < nth; ++jth) { + auto R = (const float *)(result_buffer + jth*size_thread); + auto Mj = R + Dv*rk2; + auto Sj = Mj + rk2; + R += j*Dv; + accumulate_qkv(Dv, M, S, Mj[j], Sj[j], Racc, R); + } + float norm = S > 0 ? 1/S : 1; + for (int i = 0; i < Dv; ++i) Racc[i] *= norm; + } + return true; + } + int gcd_k = simple_gcd(nstep_k, nth); + if (gcd_k >= 1) { + int nth_k = nth/gcd_k; + int ith_k = ith%gcd_k; + int ith_q = ith/gcd_k; + int nq_per_thread = (rk2 + nth_k - 1)/nth_k; + if (nq_per_thread > 1) { + int ith_mid = nth_k; + int nq_this_thread = nq_per_thread; + if (nq_per_thread*nth_k > rk2) { + ith_mid = rk2 - nth_k*(nq_per_thread - 1); + if (ith_q >= ith_mid) --nq_this_thread; + } + int j_mid = ith_mid*nq_per_thread; + auto work = (char *)work_buffer; + auto size_thread = (Dv + 16)*nq_per_thread*sizeof(float); + auto result_buffer = work; + + auto kth = (const char *)k + ith_k*(nek1/gcd_k)*stride_k; + auto vth = (const char *)v + ith_k*(nek1/gcd_k)*stride_v; + auto q_offset = ith_q < ith_mid ? ith_q*nq_per_thread*nbq2 : (ith_mid*nq_per_thread + (ith_q - ith_mid)*nq_this_thread)*nbq2; + auto qth = (const char *)q + q_offset; + auto mth = mask ? (const char *)mask + ith_k*(nek1/gcd_k)*sizeof(uint16_t) : nullptr; // we don't have ggml_half available here + + // Each thread will produce a result of size Dv*nq_this_thread*sizeof(float) + // In addition, we need M, S for the nq_this_thread rows the thread is processing + // => (Dv + 2)*nq_per_thread*sizeof(float). We use (Dv + 16) instead to make sure threads are not + // writing onto the same cache line. + auto work_this_thread = (float *)(result_buffer + ith*size_thread); + if (!iqk_flash_attn_impl(int_type_k, int_type_v, + Dk, Dv, nq_this_thread, nek1/gcd_k, nbq2, stride_k, stride_v, 0, Dv, //Dk*sizeof(uint16_t), Dv, + (const float *)qth, (const void *)kth, (const void *)vth, (const void *)mth, nullptr, 0, + scale, softcap, + work_this_thread, work_this_thread + (Dv+0)*nq_this_thread, work_this_thread + (Dv+1)*nq_this_thread)) return false; + + barrier(barrier_data); + + // There are nek1/gcd_k contributions for each j that we need to sum up + // Thread i computed k/v (i%gcd_k)*(nek1/gcd_k) for j (i/gcd_k)*(rk2/nth_k)...((i/gcd_k)+1)*(rk2/nth_k) and results at offset i*size_thread + + // TODO: simdify this + // TODO: if nth > rk2, have threads process portions of the rows instead of entire rows as it is now + for (int j = ith; j < rk2; j += nth) { + auto Racc = qkv + j*nb1/sizeof(float); + float M = -INFINITY, S = 0; + int jth_first, jj, nq_this_j; + if (j < j_mid) { + jth_first = j/nq_per_thread; + jj = j%nq_per_thread; + nq_this_j = nq_per_thread; + } else { + jth_first = ith_mid + (j - j_mid)/(nq_per_thread-1); + jj = (j - j_mid)%(nq_per_thread-1); + nq_this_j = nq_per_thread - 1; + } + jth_first *= gcd_k; + for (int jth = jth_first; jth < jth_first + gcd_k; ++jth) { + auto R = (const float *)(result_buffer + jth*size_thread); + auto Mj = R + Dv*nq_this_j; + auto Sj = Mj + nq_this_j; + R += jj*Dv; + accumulate_qkv(Dv, M, S, Mj[jj], Sj[jj], Racc, R); + } + float norm = S > 0 ? 1/S : 1; + for (int i = 0; i < Dv; ++i) Racc[i] *= norm; + } + return true; + } + } + } + + if (neq3 == 1 && rk2 > 1 && rk2 == rv2 && neq1 == 1 && nth >= 1 && nek2*nek1 >= 32*nth) { + auto result_size = (Dv + 16)*rk2*sizeof(float); + int gcd = simple_gcd(nek2, nth); + int nth_k = nth/gcd; + int nek2_k = nek2/gcd; + int nchunk = nek2_k*nek1/32; + int npt = (nchunk + nth_k - 1)/nth_k; + int nk; + if (npt*nth_k == nchunk) { + nk = 32 * (nek2*nek1/(32*nth)); + } else { + //int nm = std::max(1, npt/8); + int nm = 1; + while (true) { + if (nm*4 >= npt) break; + nm *= 2; + } + nk = 32*nm; + } + //int nk = 32 * (nek2*nek1/(32*nth)); + int nkk = (nek1 + nk - 1)/nk; + int nstep_k = nek2*nkk; + //if (ith == 0) printf("rk2 = %d, nek1 = %d, nek2 = %d, nk = %d, nkk = %d, nstep_k = %d\n", (int)rk2, (int)nek1, (int)nek2, nk, nkk, nstep_k); + for (int istep_k = ith; istep_k < nstep_k; istep_k += nth) { + int ik02 = istep_k/nkk; + int ik01 = nk*(istep_k - ik02*nkk); + int this_nk = ik01 + nk <= nek1 ? nk : nek1 - ik01; + if (this_nk <= 0) break; + auto this_result = (float *)((char *)work_buffer + istep_k*result_size); + auto this_q = (const float *)((const char *)q + ik02*rk2*nbq2); + auto this_k = (const char *)k + ik01*stride_k + ik02*nbk2; + auto this_v = (const char *)v + ik01*stride_v + ik02*nbv2; + auto this_m = mask ? (const char *)mask + ik01*sizeof(uint16_t) : nullptr; // we don't have ggml_half available here + if (!iqk_flash_attn_impl(int_type_k, int_type_v, + Dk, Dv, rk2, this_nk, nbq2, stride_k, stride_v, 0, Dv, + this_q, (const void *)this_k, (const void *)this_v, (const void *)this_m, nullptr, 0, + scale, softcap, this_result, this_result + (Dv+0)*rk2, this_result + (Dv+1)*rk2)) return false; + } + + barrier(barrier_data); + + // We have nkk results for each head + for (int iq2 = ith; iq2 < neq2; iq2 += nth) { + // ik02*rk2 + il = iq2 (il = 0...rk2-1) => ik02 = iq2/rk2, il = iq2%rk2; + int ik02 = iq2/rk2; + int il = iq2 - ik02*rk2; + auto Racc = qkv + iq2*nb1/sizeof(float); + //std::memset(Racc, 0, Dv*sizeof(float)); + float M = -INFINITY, S = 0; + for (int ikk = 0; ikk < nkk; ++ikk) { + int istep_k = ik02*nkk + ikk; + auto this_result = (float *)((char *)work_buffer + istep_k*result_size); + const float * R = this_result + il*Dv; + const float * Mj = this_result + Dv*rk2; + const float * Sj = Mj + rk2; + accumulate_qkv(Dv, M, S, Mj[il], Sj[il], Racc, R); + } + if (sinks) { + float s = ((const float *)sinks)[iq2]; + if (s > M) { + float m = expf(M - s); + for (int i = 0; i < Dv; ++i) Racc[i] *= m; + S = S*m + 1; + } else { + S += expf(s - M); + } + } + float norm = S > 0 ? 1/S : 1; + for (int i = 0; i < Dv; ++i) Racc[i] *= norm; + } + return true; + } + + // I keep changing my mind what is the best strategy to split the threads when processing + // multiple heads. This is my current thinking, the commented out code below was the previous. + int ntg = nth/simple_gcd(neq2*neq3, nth); + int neq1g = (neq1 + ntg - 1)/ntg; + //int64_t work_per_slice = D*nek1*neq1; + //int ntg = 1; + // + // When neq1 is large, it is better to have more than one thread process one (iq2,iq3) matrix + // But we also want each thread to process the same amount of rows, so neq1 must be a multiple of + // the number of threads processing the (iq2, iq3) matrix. + // + //if (neq1 >= 8*nth) { + // if (nth%8 == 0 && neq1%8 == 0 && work_per_slice >= (1 << 23)) ntg = 8; + // else if (nth%4 == 0 && neq1%4 == 0 && work_per_slice >= (1 << 21)) ntg = 4; + // else if (nth%2 == 0 && neq1%2 == 0 && work_per_slice >= (1 << 19)) ntg = 2; + //} + int counter = 0; + for (int64_t iq3 = 0; iq3 < neq3; iq3++) { + for (int64_t iq2 = 0; iq2 < neq2; iq2++) { + auto sinksf = sinks ? (const float *)sinks + iq2 : nullptr; + if (counter++ % (nth/ntg) == ith/ntg) { + int iq1 = (ith%ntg)*neq1g; + int this_neq1 = std::min(neq1g, neq1-iq1); + if (this_neq1 > 0) { + if (!iqk_flash_attn_impl(int_type_k, int_type_v, + Dk, Dv, this_neq1, nek1, stride_q, stride_k, stride_v, stride_m, ne1*nb1/sizeof(float), + (const float *)((const char *)q + iq2*nbq2 + iq3*nbq3 + iq1*stride_q), + (const void *)((const char *)k + iq2/rk2*nbk2 + iq3/rk3*nbk3), + (const void *)((const char *)v + iq2/rv2*nbv2 + iq3/rv3*nbv3), + mask ? (const void *)((const char *)mask + iq1*stride_m) : nullptr, sinksf, 1, + scale, softcap, + (float *)((char *)qkv + (iq3*ne2*ne1 + iq2 + iq1*ne1)*nb1), nullptr, nullptr)) return false; + } + } + } + } + + return true; +} + +#else + +bool iqk_flash_attn_noalibi([[maybe_unused]] int type_q, [[maybe_unused]] int type_mask, [[maybe_unused]] float max_bias, + [[maybe_unused]] int neq3, [[maybe_unused]] int neq2, [[maybe_unused]] long nbq3, [[maybe_unused]] long nbq2, + [[maybe_unused]] int nek3, [[maybe_unused]] int nek2, [[maybe_unused]] long nbk3, [[maybe_unused]] long nbk2, + [[maybe_unused]] int nev3, [[maybe_unused]] int nev2, [[maybe_unused]] long nbv3, [[maybe_unused]] long nbv2, + [[maybe_unused]] int ne2, [[maybe_unused]] int ne1, [[maybe_unused]] long nb1, + [[maybe_unused]] int type_k, // type of k + [[maybe_unused]] int type_v, // type of v + [[maybe_unused]] int Dk, // K head size + [[maybe_unused]] int Dv, // V head size + [[maybe_unused]] int nq, // number of columns in q + [[maybe_unused]] int nk, // number of rows in k + [[maybe_unused]] int stride_q, // distance between q columns in bytes + [[maybe_unused]] int stride_k, // distance between k rows in bytes + [[maybe_unused]] int stride_v, // distance between v rows in bytes + [[maybe_unused]] int stride_m, // distance between mask rows (in bytes + [[maybe_unused]] const void * q, // q matrix. + [[maybe_unused]] const void * k, // k matrix. Assumed to be fp16, nq x nk elements + [[maybe_unused]] const void * v, // v matrix. Assumed to be fp16, nq x nk elements + [[maybe_unused]] const void * mask, // mask. If not null, assumed to be fp16. nq x nk elements + [[maybe_unused]] float scale, // scale applied before softmax + [[maybe_unused]] float softcap, // if > 0, a "soft-cap" operation is applied before softmax + [[maybe_unused]] float * qkv, // v*softmax(scale*(k*q)) + [[maybe_unused]] void * work_buffer, [[maybe_unused]] barrier_t barrier, [[maybe_unused]] void * barrier_data, + [[maybe_unused]] int ith, [[maybe_unused]] int nth, [[maybe_unused]] int n_swa) { + return false; +} + +#endif + diff --git a/llama.cpp/ggml/src/iqk/iqk_flash_impl.h b/llama.cpp/ggml/src/iqk/iqk_flash_impl.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_flash_impl.h @@ -0,0 +1,35 @@ +// +// Copyright (C) 2024-2025 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include + +bool iqk_flash_attn_impl(int type_k, // type of k + int type_v, // type of v + int Dk, // K head size + int Dv, // V head size + int nq, // number of columns in q + int nk, // number of rows in k + int stride_q, // distance between q columns in bytes + int stride_k, // distance between k rows in bytes + int stride_v, // distance between v rows in bytes + int stride_m, // distance between mask rows (in bytes + int stride_qkv, // distance between rows in mask (in bytes) + const float * q, // q matrix. + const void * k, // k matrix. Assumed to be fp16, nq x nk elements + const void * v, // v matrix. Assumed to be fp16, nq x nk elements + const void * mask, // mask. If not null, assumed to be fp16. nq x nk elements + const float * sinksf, // attention sinks + int nsinks, // number of sinks + float scale, // scale applied before softmax + float softcap, // if > 0, a "soft-cap" operation is applied before softmax + float * qkv, // v*softmax(scale*(k*q)) + float * M, + float * S); + +void * iqk_repack_k(int type_k, int nek0, int nek1, int nek2, int nek3, long nbk1, long nbk2, long nbk3, + const void * k, void * work, int ith, int nth, int& repacked_type, uint64_t& row_size); diff --git a/llama.cpp/ggml/src/iqk/iqk_gemm_floats.cpp b/llama.cpp/ggml/src/iqk/iqk_gemm_floats.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_gemm_floats.cpp @@ -0,0 +1,1144 @@ +#include "iqk_gemm_floats.h" + +#ifdef IQK_IMPLEMENT + +#include "ggml-impl.h" + +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" + +#ifdef __x86_64__ + +namespace { + +// float matrices - we handle f16, bf16 (if native bf16 support is available) and f32, but only to f32 result + +struct QFBase { +#ifdef __AVX512F__ + constexpr static int k_step = 16; + using Data = __m512; + using Acc = __m512; + static inline Data load(const ggml_half * x) { return _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)x)); } + static inline Data load(const float * x) { return _mm512_loadu_ps(x); } + static inline Data load(const ggml_bf16_t * x) { + return _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256((const __m256i*)x)), 16)); + } + static inline Acc acc(Acc prev, const Data& y, const Data& x) { + return _mm512_fmadd_ps(y, x, prev); + } + static inline Acc acc_first(const Data& y, const Data& x) { + return _mm512_mul_ps(y, x); + } + static inline Acc add(Acc x, Acc y) { return _mm512_add_ps(x, y); } + static inline float hsum(Acc acc) { + return _mm512_reduce_add_ps(acc); + } + template + static inline Data load4Floats(const Float * x) { + return _mm512_insertf32x4(_mm512_setzero_ps(), load128(x), 0); + } + static inline Acc acc_r4(Acc acc, const Data * xv, const Data& yv) { + acc = _mm512_fmadd_ps(xv[0], _mm512_shuffle_ps(yv, yv, 0x00), acc); + acc = _mm512_fmadd_ps(xv[1], _mm512_shuffle_ps(yv, yv, 0x55), acc); + acc = _mm512_fmadd_ps(xv[2], _mm512_shuffle_ps(yv, yv, 0xaa), acc); + acc = _mm512_fmadd_ps(xv[3], _mm512_shuffle_ps(yv, yv, 0xff), acc); + return acc; + } + static inline Acc acc_r4_first(const Data * xv, const Data& yv) { + auto acc = _mm512_mul_ps(xv[0], _mm512_shuffle_ps(yv, yv, 0x00)); + acc = _mm512_fmadd_ps(xv[1], _mm512_shuffle_ps(yv, yv, 0x55), acc); + acc = _mm512_fmadd_ps(xv[2], _mm512_shuffle_ps(yv, yv, 0xaa), acc); + acc = _mm512_fmadd_ps(xv[3], _mm512_shuffle_ps(yv, yv, 0xff), acc); + return acc; + } + static inline __m128 hsum_r4(Acc acc) { + auto sum1 = _mm_add_ps(_mm512_extractf32x4_ps(acc, 0), _mm512_extractf32x4_ps(acc, 1)); + auto sum2 = _mm_add_ps(_mm512_extractf32x4_ps(acc, 2), _mm512_extractf32x4_ps(acc, 3)); + return _mm_add_ps(sum1, sum2); + } +#else + constexpr static int k_step = 8; + using Data = __m256; + using Acc = __m256; + static inline Data load(const ggml_half * x) { return _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)x)); } + static inline Data load(const float * x) { return _mm256_loadu_ps(x); } + static inline Data load(const ggml_bf16_t * x) { + return _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(_mm_loadu_si128((const __m128i*)x)), 16)); + } + static inline Acc acc(Acc prev, const Data& y, const Data& x) { + return _mm256_fmadd_ps(y, x, prev); + } + static inline Acc add(Acc x, Acc y) { return _mm256_add_ps(x, y); } + static inline Acc acc_r4(Acc acc, const Data * xv, const Data& yv) { + acc = _mm256_fmadd_ps(xv[0], _mm256_shuffle_ps(yv, yv, 0x00), acc); + acc = _mm256_fmadd_ps(xv[1], _mm256_shuffle_ps(yv, yv, 0x55), acc); + acc = _mm256_fmadd_ps(xv[2], _mm256_shuffle_ps(yv, yv, 0xaa), acc); + acc = _mm256_fmadd_ps(xv[3], _mm256_shuffle_ps(yv, yv, 0xff), acc); + return acc; + } + static inline Acc acc_r4_first(const Data * xv, const Data& yv) { + auto acc = _mm256_mul_ps(xv[0], _mm256_shuffle_ps(yv, yv, 0x00)); + acc = _mm256_fmadd_ps(xv[1], _mm256_shuffle_ps(yv, yv, 0x55), acc); + acc = _mm256_fmadd_ps(xv[2], _mm256_shuffle_ps(yv, yv, 0xaa), acc); + acc = _mm256_fmadd_ps(xv[3], _mm256_shuffle_ps(yv, yv, 0xff), acc); + return acc; + } + static inline Acc acc_first(const Data& y, const Data& x) { + return _mm256_mul_ps(y, x); + } + static inline float hsum(Acc acc) { + return hsum_float_8(acc); + } + static inline __m128 hsum_r4(Acc acc) { + return _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1)); + } + template + static inline Data load4Floats(const Float * x) { + return _mm256_insertf128_ps(_mm256_setzero_ps(), load128(x), 0); + } +#endif + static inline __m128 load128(const ggml_half * x) { return _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)x)); } + static inline __m128 load128(const float * x) { return _mm_loadu_ps(x); } + static inline __m128 load128(const ggml_bf16_t * x) { + return _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i*)x)), 16)); + } +}; + +template struct QFT final : public QFBase { + constexpr static int nrc = nrc_in; + QFT(const DataInfo& info) { + for (int iy = 0; iy < nrc; ++iy) y[iy] = (const Float *)info.src1_row(iy); + } + QFT(const char * cx, size_t bx) { + for (int iy = 0; iy < nrc; ++iy) y[iy] = (const Float *)(cx + iy*bx); + } + IQK_ALWAYS_INLINE Data load1(int iy, int i) const { return load(y[iy] + k_step*i); } + IQK_ALWAYS_INLINE Data load_tail(int iy, int i) const { return load4Floats(y[iy] + 4*i); } + IQK_ALWAYS_INLINE void load_r4(int ix, int i, Data * xv) const { + xv[0] = load1(ix+0, i); + xv[1] = load1(ix+1, i); + xv[2] = load1(ix+2, i); + xv[3] = load1(ix+3, i); +#ifdef __AVX512F__ + auto t0 = _mm512_unpacklo_ps(xv[0], xv[1]); + auto t1 = _mm512_unpacklo_ps(xv[2], xv[3]); + auto t2 = _mm512_unpackhi_ps(xv[0], xv[1]); + auto t3 = _mm512_unpackhi_ps(xv[2], xv[3]); + xv[0] = _mm512_castpd_ps(_mm512_unpacklo_pd(_mm512_castps_pd(t0), _mm512_castps_pd(t1))); + xv[1] = _mm512_castpd_ps(_mm512_unpackhi_pd(_mm512_castps_pd(t0), _mm512_castps_pd(t1))); + xv[2] = _mm512_castpd_ps(_mm512_unpacklo_pd(_mm512_castps_pd(t2), _mm512_castps_pd(t3))); + xv[3] = _mm512_castpd_ps(_mm512_unpackhi_pd(_mm512_castps_pd(t2), _mm512_castps_pd(t3))); +#else + auto t0 = _mm256_unpacklo_ps(xv[0], xv[1]); + auto t1 = _mm256_unpacklo_ps(xv[2], xv[3]); + auto t2 = _mm256_unpackhi_ps(xv[0], xv[1]); + auto t3 = _mm256_unpackhi_ps(xv[2], xv[3]); + xv[0] = _mm256_castpd_ps(_mm256_unpacklo_pd(_mm256_castps_pd(t0), _mm256_castps_pd(t1))); + xv[1] = _mm256_castpd_ps(_mm256_unpackhi_pd(_mm256_castps_pd(t0), _mm256_castps_pd(t1))); + xv[2] = _mm256_castpd_ps(_mm256_unpacklo_pd(_mm256_castps_pd(t2), _mm256_castps_pd(t3))); + xv[3] = _mm256_castpd_ps(_mm256_unpackhi_pd(_mm256_castps_pd(t2), _mm256_castps_pd(t3))); +#endif + } + const Float * y[nrc]; +}; + +// TBD if we want this +//template +//IQK_NOINLINE void mul_mat_Qx_Qy_Mx1(int n, const char * cx, size_t bx, int ix0, const DataInfo& info) { +// static_assert(Qy::nrc == 1); +// int nb = n/QFBase::k_step; +// int nb4 = n/4; +// Qy y(info); +// Qx x(cx + ix0*bx, bx); +// QFBase::Data xv[2*Qx::nrc]; +// QFBase::Acc acc[2*Qx::nrc]; +// auto yv1 = y.load1(0, 0); +// auto yv2 = y.load1(0, 1); +// for (int ix = 0; ix < Qx::nrc; ++ix) { +// xv[2*ix+0] = x.load1(ix, 0); +// xv[2*ix+1] = x.load1(ix, 1); +// acc[2*ix+0] = QFBase::acc_first(yv1, xv[2*ix+0]); +// acc[2*ix+1] = QFBase::acc_first(yv2, xv[2*ix+1]); +// } +// for (int i = 1; i < nb/2; ++i) { +// yv1 = y.load1(0, 2*i+0); +// yv2 = y.load1(0, 2*i+1); +// for (int ix = 0; ix < Qx::nrc; ++ix) { +// xv[2*ix+0] = x.load1(ix, 2*i+0); +// xv[2*ix+1] = x.load1(ix, 2*i+1); +// acc[2*ix+0] = QFBase::acc(acc[2*ix+0], yv1, xv[2*ix+0]); +// acc[2*ix+1] = QFBase::acc(acc[2*ix+1], yv2, xv[2*ix+1]); +// } +// } +// for (int i = (QFBase::k_step/4)*nb; i < nb4; ++i) { +// yv1 = y.load_tail(0, i); +// for (int ix = 0; ix < Qx::nrc; ++ix) { +// xv[ix] = x.load_tail(ix, i); +// acc[2*ix+0] = QFBase::acc(acc[2*ix+0], yv1, xv[ix]); +// } +// } +// for (int ix = 0; ix < Qx::nrc; ++ix) info.store(ix0+ix, 0, QFBase::hsum(QFBase::add(acc[2*ix+0], acc[2*ix+1]))); +//} + +template +IQK_NOINLINE void mul_mat_Qx_Qy_MxN(int n, const char * cx, size_t bx, int ix0, const DataInfo& info) { + int nb = n/QFBase::k_step; + int nb4 = n/4; + Qy y(info); + Qx x(cx + ix0*bx, bx); + QFBase::Data xv[Qx::nrc]; + QFBase::Acc acc[Qx::nrc*Qy::nrc]; + auto yv = y.load1(0, 0); + for (int ix = 0; ix < Qx::nrc; ++ix) { + xv[ix] = x.load1(ix, 0); + acc[ix] = QFBase::acc_first(yv, xv[ix]); + } + for (int iy = 1; iy < Qy::nrc; ++iy) { + yv = y.load1(iy, 0); + for (int ix = 0; ix < Qx::nrc; ++ix) acc[Qx::nrc*iy + ix] = QFBase::acc_first(yv, xv[ix]); + } + for (int i = 1; i < nb; ++i) { + yv = y.load1(0, i); + for (int ix = 0; ix < Qx::nrc; ++ix) { + xv[ix] = x.load1(ix, i); + acc[ix] = QFBase::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < Qy::nrc; ++iy) { + yv = y.load1(iy, i); + for (int ix = 0; ix < Qx::nrc; ++ix) acc[Qx::nrc*iy + ix] = QFBase::acc(acc[Qx::nrc*iy + ix], yv, xv[ix]); + } + } + for (int i = (QFBase::k_step/4)*nb; i < nb4; ++i) { + yv = y.load_tail(0, i); + for (int ix = 0; ix < Qx::nrc; ++ix) { + xv[ix] = x.load_tail(ix, i); + acc[ix] = QFBase::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < Qy::nrc; ++iy) { + yv = y.load_tail(iy, i); + for (int ix = 0; ix < Qx::nrc; ++ix) acc[Qx::nrc*iy + ix] = QFBase::acc(acc[Qx::nrc*iy + ix], yv, xv[ix]); + } + } + for (int iy = 0; iy < Qy::nrc; ++iy) for (int ix = 0; ix < Qx::nrc; ++ix) info.store(ix0+ix, iy, QFBase::hsum(acc[Qx::nrc*iy+ix])); +} + +template +inline void mul_mat_Qx_Qy_MxN_fa(int n, const char * cx, size_t bx, int ix0, const DataInfo& info) { + int nb = n/QFBase::k_step; + Qy y(info); + Qx x(cx + ix0*bx, bx); + QFBase::Data xv[Qx::nrc]; + QFBase::Acc acc[Qx::nrc*Qy::nrc]; + auto yv = y.load1(0, 0); + for (int ix = 0; ix < Qx::nrc; ++ix) { + xv[ix] = x.load1(ix, 0); + acc[ix] = QFBase::acc_first(yv, xv[ix]); + } + for (int iy = 1; iy < Qy::nrc; ++iy) { + yv = y.load1(iy, 0); + for (int ix = 0; ix < Qx::nrc; ++ix) acc[Qx::nrc*iy + ix] = QFBase::acc_first(yv, xv[ix]); + } + for (int i = 1; i < nb; ++i) { + yv = y.load1(0, i); + for (int ix = 0; ix < Qx::nrc; ++ix) { + xv[ix] = x.load1(ix, i); + acc[ix] = QFBase::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < Qy::nrc; ++iy) { + yv = y.load1(iy, i); + for (int ix = 0; ix < Qx::nrc; ++ix) acc[Qx::nrc*iy + ix] = QFBase::acc(acc[Qx::nrc*iy + ix], yv, xv[ix]); + } + } + for (int iy = 0; iy < Qy::nrc; ++iy) for (int ix = 0; ix < Qx::nrc; ++ix) info.store(ix0+ix, iy, QFBase::hsum(acc[Qx::nrc*iy+ix])); +} + +template +inline void mul_mat_Qx_Qy_MxN_fa4(int D, const char * cx, size_t bx, int ix0, const DataInfo& info) { + static_assert(Qx::nrc%4 == 0); + int nb = D/QFBase::k_step; + Qy y(info); + Qx x(cx + ix0*bx, bx); + QFBase::Data xv[Qx::nrc]; + QFBase::Acc acc[Qx::nrc*Qy::nrc/4] = {}; + for (int i = 0; i < nb; ++i) { + for (int ix = 0; ix < Qx::nrc/4; ++ix) x.load_r4(4*ix, i, xv + 4*ix); + for (int iy = 0; iy < Qy::nrc; ++iy) { + auto yv = y.load1(iy, i); + for (int ix = 0; ix < Qx::nrc/4; ++ix) acc[ix*Qy::nrc + iy] = QFBase::acc_r4(acc[ix*Qy::nrc + iy], xv + 4*ix, yv); + } + } + for (int iy = 0; iy < Qy::nrc; ++iy) { + for (int ix = 0; ix < Qx::nrc/4; ++ix) info.store(ix0+4*ix, iy, QFBase::hsum_r4(acc[ix*Qy::nrc + iy])); + } +} + +// This will handle any of f16 x f32, f32 x f16, f16 x f16, f32 x f32, with computations done +// in f32 (i.e., f16 is first converted to f32). It is easy to extend to computations done in +// f16, but I don't have a CPU capable of f16 vector arithmetic, so not doing it for now. +template +void mul_mat_fX_fY_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + const char * cx = (const char *)vx; + // TBD if we want this + //if constexpr (nrc_y == 1) { + // constexpr int k_nx = 2; + // for (int ix = 0; ix < nrc_x/k_nx; ++ix) { + // mul_mat_Qx_Qy_Mx1, QFT>(n, cx, bx, ix*k_nx, info); + // } + // if (int lastx = k_nx*(nrc_x/k_nx); lastx < nrc_x) { + // int nx = nrc_x - lastx; + // switch (nx) { + // case 1: mul_mat_Qx_Qy_Mx1, QFT>(n, cx, bx, lastx, info); break; + // case 2: mul_mat_Qx_Qy_Mx1, QFT>(n, cx, bx, lastx, info); break; + // case 3: mul_mat_Qx_Qy_Mx1, QFT>(n, cx, bx, lastx, info); break; + // } + // //mul_mat_Qx_Qy_Mx1, QFT>(n, cx, bx, lastx, info); + // } + // return; + //} +#ifdef __AVX512F__ + constexpr int k_nx = 5; +#else + constexpr int k_nx = nrc_y == 1 ? 4 : 2; +#endif + for (int ix = 0; ix < nrc_x/k_nx; ++ix) { + mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, ix*k_nx, info); + } + int last_x = k_nx*(nrc_x/k_nx); + if (last_x == nrc_x) return; + int nx = nrc_x - last_x; +#ifdef __AVX512F__ + switch (nx) { + case 1: mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, last_x, info); break; + case 2: mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, last_x, info); break; + case 3: mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, last_x, info); break; + case 4: mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, last_x, info); break; + } +#else + if constexpr (nrc_y == 1) { + switch (nx) { + case 1: mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, last_x, info); break; + case 2: mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, last_x, info); break; + case 3: mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, last_x, info); break; + } + } else { + switch (nx) { + case 1: mul_mat_Qx_Qy_MxN, QFT>(n, cx, bx, last_x, info); break; + } + } +#endif +} + +#ifdef __AVX512BF16__ +template +static void mul_mat_bf16_r16_bf16(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%16 == 0); + const ggml_bf16_t * y[nrc_y]; + static_for([&](const int iy) { y[iy] = (const ggml_bf16_t *)info.src1_row(iy); }); + + for (int ix = 0; ix < nrc_x/32; ++ix) { + __m512 acc[2*nrc_y] = {}; + __m512bh qx[8]; + const ggml_bf16_t * b8_1 = (const ggml_bf16_t *)((const char *)vx + (32*ix+ 0)*bx); + const ggml_bf16_t * b8_2 = (const ggml_bf16_t *)((const char *)vx + (32*ix+16)*bx); + for (int ib = 0; ib < n/8; ++ib) { + qx[0] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8_1+4*ib+0); + qx[1] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8_1+4*ib+1); + qx[2] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8_1+4*ib+2); + qx[3] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8_1+4*ib+3); + qx[4] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8_2+4*ib+0); + qx[5] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8_2+4*ib+1); + qx[6] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8_2+4*ib+2); + qx[7] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8_2+4*ib+3); + static_for([&](const int iy) { + auto y128 = _mm_loadu_si128((const __m128i*)y[iy]+ib); + //auto y = _mm512_broadcast_i32x4(y128); + auto y256 = MM256_SET_M128I(y128, y128); + auto y = _mm512_inserti32x8(_mm512_castsi256_si512(y256), y256, 1); + acc[2*iy+0] = _mm512_dpbf16_ps(acc[2*iy+0], qx[0], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + acc[2*iy+0] = _mm512_dpbf16_ps(acc[2*iy+0], qx[1], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + acc[2*iy+0] = _mm512_dpbf16_ps(acc[2*iy+0], qx[2], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + acc[2*iy+0] = _mm512_dpbf16_ps(acc[2*iy+0], qx[3], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + acc[2*iy+1] = _mm512_dpbf16_ps(acc[2*iy+1], qx[4], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + acc[2*iy+1] = _mm512_dpbf16_ps(acc[2*iy+1], qx[5], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + acc[2*iy+1] = _mm512_dpbf16_ps(acc[2*iy+1], qx[6], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + acc[2*iy+1] = _mm512_dpbf16_ps(acc[2*iy+1], qx[7], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + }); + } + static_for([&](const int iy) { + info.store(32*ix+ 0, iy, acc[2*iy+0]); + info.store(32*ix+16, iy, acc[2*iy+1]); + }); + } + for (int ix = 32*(nrc_x/32); ix < nrc_x; ix += 16) { + __m512 acc[nrc_y] = {}; + __m512bh qx[4]; + const ggml_bf16_t * b8 = (const ggml_bf16_t *)((const char *)vx + (ix+0)*bx); + for (int ib = 0; ib < n/8; ++ib) { + qx[0] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8+4*ib+0); + qx[1] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8+4*ib+1); + qx[2] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8+4*ib+2); + qx[3] = (__m512bh)_mm512_loadu_si512((const __m512i *)b8+4*ib+3); + static_for([&](const int iy) { + auto y128 = _mm_loadu_si128((const __m128i*)y[iy]+ib); + auto y256 = MM256_SET_M128I(y128, y128); + auto y = _mm512_inserti32x8(_mm512_castsi256_si512(y256), y256, 1); + acc[iy] = _mm512_dpbf16_ps(acc[iy], qx[0], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + acc[iy] = _mm512_dpbf16_ps(acc[iy], qx[1], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + acc[iy] = _mm512_dpbf16_ps(acc[iy], qx[2], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + acc[iy] = _mm512_dpbf16_ps(acc[iy], qx[3], (__m512bh)_mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + }); + } + static_for([&](const int iy) { + info.store(ix, iy, acc[iy]); + }); + } +} + +struct QFBaseBF16 { + constexpr static int k_step = 32; + using Data = __m512bh; + using Acc = __m512; + static inline Data load(const ggml_bf16_t * x) { return __m512bh(_mm512_loadu_si512((const __m512i *)x)); } + static inline Acc acc(Acc prev, Data y, Data x) { + return _mm512_dpbf16_ps(prev, y, x); + } + static inline Acc acc_first(const Data& y, const Data& x) { + return _mm512_dpbf16_ps(_mm512_setzero_ps(), y, x); + } + static inline float hsum(Acc acc) { + return _mm512_reduce_add_ps(acc); + } +}; +template struct QFTBF16 final : public QFBaseBF16 { + constexpr static int nrc = nrc_in; + QFTBF16(const DataInfo& info) { + for (int iy = 0; iy < nrc; ++iy) y[iy] = (const ggml_bf16_t *)info.src1_row(iy); + } + QFTBF16(const char * cx, size_t bx) { + for (int iy = 0; iy < nrc; ++iy) y[iy] = (const ggml_bf16_t *)(cx + iy*bx); + } + IQK_ALWAYS_INLINE Data load1(int iy, int i) const { return load(y[iy] + k_step*i); } + const ggml_bf16_t * y[nrc]; +}; +struct QFBaseBF16x8 { + constexpr static int k_step = 16; + using Data = __m256bh; + using Acc = __m256; + static inline Data load(const ggml_bf16_t * x) { return __m256bh(_mm256_loadu_si256((const __m256i *)x)); } + static inline Acc acc(Acc prev, Data y, Data x) { + return _mm256_dpbf16_ps(prev, y, x); + } + static inline Acc acc_first(const Data& y, const Data& x) { + return _mm256_dpbf16_ps(_mm256_setzero_ps(), y, x); + } + static inline float hsum(Acc acc) { + return hsum_float_8(acc); + } +}; +template struct QFTBF16x8 final : public QFBaseBF16x8 { + constexpr static int nrc = nrc_in; + QFTBF16x8(const DataInfo& info) { + for (int iy = 0; iy < nrc; ++iy) y[iy] = (const ggml_bf16_t *)info.src1_row(iy); + } + QFTBF16x8(const char * cx, size_t bx) { + for (int iy = 0; iy < nrc; ++iy) y[iy] = (const ggml_bf16_t *)(cx + iy*bx); + } + IQK_ALWAYS_INLINE Data load1(int iy, int i) const { return load(y[iy] + k_step*i); } + const ggml_bf16_t * y[nrc]; +}; + +template +IQK_NOINLINE void mul_mat_Qx_Qy_MxN(int n, const char * cx, size_t bx, int ix0, const DataInfo& info) { + int nb = n/QFBaseBF16::k_step; + QFTBF16 y(info); + QFTBF16 x(cx + ix0*bx, bx); + QFBaseBF16::Data xv[nrc_x]; + QFBaseBF16::Acc acc[nrc_x*nrc_y]; + auto yv = y.load1(0, 0); + for (int ix = 0; ix < nrc_x; ++ix) { + xv[ix] = x.load1(ix, 0); + acc[ix] = QFBaseBF16::acc_first(yv, xv[ix]); + } + for (int iy = 1; iy < nrc_y; ++iy) { + yv = y.load1(iy, 0); + for (int ix = 0; ix < nrc_x; ++ix) acc[nrc_x*iy + ix] = QFBaseBF16::acc_first(yv, xv[ix]); + } + for (int i = 1; i < nb; ++i) { + yv = y.load1(0, i); + for (int ix = 0; ix < nrc_x; ++ix) { + xv[ix] = x.load1(ix, i); + acc[ix] = QFBaseBF16::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < nrc_y; ++iy) { + yv = y.load1(iy, i); + for (int ix = 0; ix < nrc_x; ++ix) acc[nrc_x*iy + ix] = QFBaseBF16::acc(acc[nrc_x*iy + ix], yv, xv[ix]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) for (int ix = 0; ix < nrc_x; ++ix) info.store(ix0+ix, iy, QFBaseBF16::hsum(acc[nrc_x*iy+ix])); +} + +template +void mul_mat_fX_fY_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + constexpr int k_nx = nrc_y <= 2 ? 8 : 5; + const char * cx = (const char *)vx; + for (int ix = 0; ix < nrc_x/k_nx; ++ix) { + mul_mat_Qx_Qy_MxN(n, cx, bx, ix*k_nx, info); + } + int last_x = k_nx*(nrc_x/k_nx); + if (last_x == nrc_x) return; + int nx = nrc_x - last_x; + if constexpr (nrc_y <= 2) { + if (nx >= 4) { + mul_mat_Qx_Qy_MxN(n, cx, bx, last_x, info); + last_x += 4; + if (last_x == nrc_x) return; + nx = nrc_x - last_x; + } + } + switch (nx) { + case 1: mul_mat_Qx_Qy_MxN(n, cx, bx, last_x, info); break; + case 2: mul_mat_Qx_Qy_MxN(n, cx, bx, last_x, info); break; + case 3: mul_mat_Qx_Qy_MxN(n, cx, bx, last_x, info); break; + case 4: mul_mat_Qx_Qy_MxN(n, cx, bx, last_x, info); break; + } +} +template +IQK_NOINLINE void mul_mat_Qx_Qy_MxNx8(int n, const char * cx, size_t bx, int ix0, const DataInfo& info) { + int nb = n/QFBaseBF16x8::k_step; + QFTBF16x8 y(info); + QFTBF16x8 x(cx + ix0*bx, bx); + QFBaseBF16x8::Data xv[nrc_x]; + QFBaseBF16x8::Acc acc[nrc_x*nrc_y]; + auto yv = y.load1(0, 0); + for (int ix = 0; ix < nrc_x; ++ix) { + xv[ix] = x.load1(ix, 0); + acc[ix] = QFBaseBF16x8::acc_first(yv, xv[ix]); + } + for (int iy = 1; iy < nrc_y; ++iy) { + yv = y.load1(iy, 0); + for (int ix = 0; ix < nrc_x; ++ix) acc[nrc_x*iy + ix] = QFBaseBF16x8::acc_first(yv, xv[ix]); + } + for (int i = 1; i < nb; ++i) { + yv = y.load1(0, i); + for (int ix = 0; ix < nrc_x; ++ix) { + xv[ix] = x.load1(ix, i); + acc[ix] = QFBaseBF16x8::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < nrc_y; ++iy) { + yv = y.load1(iy, i); + for (int ix = 0; ix < nrc_x; ++ix) acc[nrc_x*iy + ix] = QFBaseBF16x8::acc(acc[nrc_x*iy + ix], yv, xv[ix]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) for (int ix = 0; ix < nrc_x; ++ix) info.store(ix0+ix, iy, QFBaseBF16x8::hsum(acc[nrc_x*iy+ix])); +} + +template +void mul_mat_fX_fY_Tx8(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + constexpr int k_nx = nrc_y <= 2 ? 8 : 5; + const char * cx = (const char *)vx; + for (int ix = 0; ix < nrc_x/k_nx; ++ix) { + mul_mat_Qx_Qy_MxNx8(n, cx, bx, ix*k_nx, info); + } + int last_x = k_nx*(nrc_x/k_nx); + if (last_x == nrc_x) return; + int nx = nrc_x - last_x; + if constexpr (nrc_y <= 2) { + if (nx >= 4) { + mul_mat_Qx_Qy_MxNx8(n, cx, bx, last_x, info); + last_x += 4; + if (last_x == nrc_x) return; + nx = nrc_x - last_x; + } + } + switch (nx) { + case 1: mul_mat_Qx_Qy_MxNx8(n, cx, bx, last_x, info); break; + case 2: mul_mat_Qx_Qy_MxNx8(n, cx, bx, last_x, info); break; + case 3: mul_mat_Qx_Qy_MxNx8(n, cx, bx, last_x, info); break; + case 4: mul_mat_Qx_Qy_MxNx8(n, cx, bx, last_x, info); break; + } +} +#endif + + +template +void set_mul_mat_f(std::array& funcs) { + for (auto& f : funcs) f = nullptr; + funcs[0] = mul_mat_fX_fY_T<1, FloatX, FloatY>; + funcs[1] = mul_mat_fX_fY_T<2, FloatX, FloatY>; + funcs[2] = mul_mat_fX_fY_T<3, FloatX, FloatY>; + funcs[3] = mul_mat_fX_fY_T<4, FloatX, FloatY>; + funcs[4] = mul_mat_fX_fY_T<5, FloatX, FloatY>; +#ifndef __AVX512F__ + funcs[5] = mul_mat_fX_fY_T<6, FloatX, FloatY>; +#endif +} + +#ifdef __AVX512BF16__ +void set_mul_mat_bf16(std::array& funcs) { + for (auto& f : funcs) f = nullptr; + funcs[0] = mul_mat_fX_fY_T<1>; + funcs[1] = mul_mat_fX_fY_T<2>; + funcs[2] = mul_mat_fX_fY_T<3>; + funcs[3] = mul_mat_fX_fY_T<4>; + funcs[4] = mul_mat_fX_fY_T<5>; +} +void set_mul_mat_bf16x8(std::array& funcs) { + for (auto& f : funcs) f = nullptr; + funcs[0] = mul_mat_fX_fY_Tx8<1>; + funcs[1] = mul_mat_fX_fY_Tx8<2>; + funcs[2] = mul_mat_fX_fY_Tx8<3>; + funcs[3] = mul_mat_fX_fY_Tx8<4>; + funcs[4] = mul_mat_fX_fY_Tx8<5>; +} +void set_mul_mat_bf16_r16(std::array& funcs) { + for (auto& f : funcs) f = nullptr; + funcs[0] = mul_mat_bf16_r16_bf16<1>; + funcs[1] = mul_mat_bf16_r16_bf16<2>; + funcs[2] = mul_mat_bf16_r16_bf16<3>; + funcs[3] = mul_mat_bf16_r16_bf16<4>; + funcs[4] = mul_mat_bf16_r16_bf16<5>; + funcs[5] = mul_mat_bf16_r16_bf16<6>; + funcs[6] = mul_mat_bf16_r16_bf16<7>; + funcs[7] = mul_mat_bf16_r16_bf16<8>; +} +#endif + +} // namespace + +bool iqk_set_kernels_float(int ne00, int typeA, int typeB, std::array& kernels) { + + if (typeA == GGML_TYPE_BF16) { + if (ne00 % 8) return false; + switch (typeB) { +#ifdef __AVX512BF16__ + case GGML_TYPE_BF16: { + if (ne00 % 16 == 0) { + set_mul_mat_bf16(kernels); + } else { + set_mul_mat_bf16x8(kernels); + } + } break; +#else + case GGML_TYPE_BF16: set_mul_mat_f(kernels); break; + case GGML_TYPE_F32: set_mul_mat_f(kernels); break; +#endif + default: return false; + } + return true; + } + + if (typeA == GGML_TYPE_BF16_R16) { + if (ne00 % 16) return false; + switch (typeB) { +#ifdef __AVX512BF16__ + case GGML_TYPE_BF16: set_mul_mat_bf16_r16(kernels); break; +#endif + default: return false; + } + return true; + } + + if (typeA == GGML_TYPE_F16 || typeA == GGML_TYPE_F32) { + if (ne00 % 4) return false; + } + if (typeA == GGML_TYPE_F16) { + switch (typeB) { + case GGML_TYPE_F16: set_mul_mat_f(kernels); break; + case GGML_TYPE_F32: set_mul_mat_f(kernels); break; + default: return false; + } + return true; + } + if (typeA == GGML_TYPE_F32) { + switch (typeB) { + case GGML_TYPE_F16: set_mul_mat_f(kernels); break; + case GGML_TYPE_F32: set_mul_mat_f(kernels); break; + default: return false; + } + return true; + } + + return false; + +} + +void iqk_gemm_default_floats(int D, int nq, const char * cx, size_t bx, DataInfo& info, int k_step) { + using q_float = float; +#ifdef HAVE_FANCY_SIMD + constexpr int nrc_q = 8; + constexpr int nrc_k = 8; +#else + // somewhat surprisingly, nrc_q = 4, nrc_k = 8 is better than nrc_q = 8, nrc_k = 4 + constexpr int nrc_q = 4; + constexpr int nrc_k = 8; +#endif + GGML_ASSERT(k_step%nrc_k == 0); + int qrem = nq - nrc_q*(nq/nrc_q); + for (int iq = 0; iq < nq/nrc_q; ++iq) { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_Qx_Qy_MxN_fa4, QFT>(D, cx, bx, ik*nrc_k, info); + } + info.cur_y += nrc_q; + } + if (qrem > 0) { + switch (qrem) { + case 1: { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_Qx_Qy_MxN_fa4, QFT>(D, cx, bx, ik*nrc_k, info); + } + } break; + case 2: { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_Qx_Qy_MxN_fa4, QFT>(D, cx, bx, ik*nrc_k, info); + } + } break; + case 3: { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_Qx_Qy_MxN_fa4, QFT>(D, cx, bx, ik*nrc_k, info); + } + } break; +#ifdef HAVE_FANCY_SIMD + case 4: { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_Qx_Qy_MxN_fa4, QFT>(D, cx, bx, ik*nrc_k, info); + } + } break; + case 5: { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_Qx_Qy_MxN_fa4, QFT>(D, cx, bx, ik*nrc_k, info); + } + } break; + case 6: { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_Qx_Qy_MxN_fa4, QFT>(D, cx, bx, ik*nrc_k, info); + } + } break; + case 7: { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_Qx_Qy_MxN_fa4, QFT>(D, cx, bx, ik*nrc_k, info); + } + } break; +#endif + } + } +} + +#else +// ----------------------------------- __aarch64__ ----------------------------------------------- + +namespace { + +struct QF16Base { + constexpr static int k_step = 8; + using Data = float16x8_t; + using Acc = float16x8_t; + static inline Data load(const __fp16 * x) { return vld1q_f16(x); } + static inline Data load4(const __fp16 * x) { return vcombine_f16(vld1_f16(x), vdup_n_f16(0)); } + static inline Acc acc(Acc prev, const Data& y, const Data& x) { + return vfmaq_f16(prev, y, x); + } + static inline Acc acc_first(const Data& y, const Data& x) { + return vmulq_f16(y, x); + } + //constexpr static int k_step = 16; + //using Data = float16x8x2_t; + //static inline Data load(const __fp16 * x) { return vld1q_f16_x2(x); } + //static inline Acc acc(Acc prev, const Data& y, const Data& x) { + // return vfmaq_f16(vfmaq_f16(prev, y.val[0], x.val[0]), y.val[1], x.val[1]); + //} + //static inline Acc acc_first(const Data& y, const Data& x) { + // return vfmaq_f16(vmulq_f16(y.val[0], x.val[0]), y.val[1], x.val[1]); + //} + static inline float hsum(Acc acc) { + float32x4_t sum = vcvt_f32_f16(vadd_f16(vget_low_f16(acc), vget_high_f16(acc))); + return vaddvq_f32(sum); + } +}; +template struct QF16 final : public QF16Base { + using Base = QF16Base; + constexpr static int nrc_y = nrc; + QF16(const DataInfo& info) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const __fp16 *)info.src1_row(iy); + } + QF16(const char * cx, size_t bx) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const __fp16 *)(cx + iy*bx); + } + IQK_ALWAYS_INLINE Data load1(int iy, int i) const { return load(y[iy] + k_step*i); } + IQK_ALWAYS_INLINE Data load_tail(int iy, int i) const { return load4(y[iy] + 4*i); } + IQK_ALWAYS_INLINE float16x8x4_t loadx(int iy, int i) const { return vld1q_f16_x4(y[iy] + 4*k_step*i); } + const __fp16 * y[nrc_y]; +}; + +struct QBF16Base { + constexpr static int k_step = 4; + using Data = float32x4_t; + using Acc = float32x4_t; + static inline Data load(const uint16_t * x) { return vreinterpretq_f32_u32(vshlq_n_u32(vmovl_u16(vld1_u16(x)), 16)); } + static inline Data load4(const uint16_t * x) { return load(x); } + static inline Acc acc(Acc prev, const Data& y, const Data& x) { + return vfmaq_f32(prev, y, x); + } + static inline Acc acc_first(const Data& y, const Data& x) { + return vmulq_f32(y, x); + } + static inline float hsum(Acc acc) { return vaddvq_f32(acc); } +}; +template struct QBF16 final : public QBF16Base { + using Base = QBF16Base; + constexpr static int nrc_y = nrc; + QBF16(const DataInfo& info) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const uint16_t *)info.src1_row(iy); + } + QBF16(const char * cx, size_t bx) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const uint16_t *)(cx + iy*bx); + } + IQK_ALWAYS_INLINE Data load1(int iy, int i) const { return load(y[iy] + k_step*i); } + IQK_ALWAYS_INLINE Data load_tail(int iy, int i) const { return load(y[iy] + 4*i); } + const uint16_t * y[nrc_y]; +}; + +struct QF32Base { + constexpr static int k_step = 4; + using Data = float32x4_t; + using Acc = float32x4_t; + static inline Data load(const float * x) { return vld1q_f32(x); } + static inline Data load4(const float * x) { return load(x); } + static inline Acc acc(Acc prev, const Data& y, const Data& x) { return vfmaq_f32(prev, y, x); } + static inline Acc acc_first(const Data& y, const Data& x) { return vmulq_f32(y, x); } + static inline float hsum(Acc acc) { return vaddvq_f32(acc); } +}; +template struct QF32 final : public QF32Base { + using Base = QF32Base; + constexpr static int nrc_y = nrc; + QF32(const DataInfo& info) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const float *)info.src1_row(iy); + } + QF32(const char * cx, size_t bx) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const float *)(cx + iy*bx); + } + IQK_ALWAYS_INLINE Data load1(int iy, int i) const { return load(y[iy] + k_step*i); } + IQK_ALWAYS_INLINE Data load_tail(int iy, int i) const { return load(y[iy] + 4*i); } + const float * y[nrc_y]; +}; + +template +IQK_NOINLINE void mul_mat_Qx_Qy_NxN(int n, const char * cx, size_t bx, int ix0, const DataInfo& info) { + GGML_ASSERT(Qx::Base::k_step == Qy::Base::k_step); + int nb = n/Qx::Base::k_step; + Qy y(info); + Qx x(cx + ix0*bx, bx); + typename Qx::Base::Data xv[Qx::nrc_y]; + typename Qx::Base::Acc acc[Qx::nrc_y*Qy::nrc_y]; + auto yv = y.load1(0, 0); + for (int ix = 0; ix < Qx::nrc_y; ++ix) { + xv[ix] = x.load1(ix, 0); + acc[ix] = Qx::Base::acc_first(yv, xv[ix]); + } + for (int iy = 1; iy < Qy::nrc_y; ++iy) { + yv = y.load1(iy, 0); + for (int ix = 0; ix < Qx::nrc_y; ++ix) acc[Qx::nrc_y*iy + ix] = Qx::Base::acc_first(yv, xv[ix]); + } + for (int i = 1; i < nb; ++i) { + yv = y.load1(0, i); + for (int ix = 0; ix < Qx::nrc_y; ++ix) { + xv[ix] = x.load1(ix, i); + acc[ix] = Qx::Base::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < Qy::nrc_y; ++iy) { + yv = y.load1(iy, i); + for (int ix = 0; ix < Qx::nrc_y; ++ix) acc[Qx::nrc_y*iy + ix] = Qx::Base::acc(acc[Qx::nrc_y*iy + ix], yv, xv[ix]); + } + } + if constexpr (Qx::Base::k_step > 4 && !is_multiple_of_k_step) { + int nb4 = n/4; + for (int i = (Qx::Base::k_step/4)*nb; i < nb4; ++i) { + yv = y.load_tail(0, i); + for (int ix = 0; ix < Qx::nrc_y; ++ix) { + xv[ix] = x.load_tail(ix, i); + acc[ix] = Qx::Base::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < Qy::nrc_y; ++iy) { + yv = y.load_tail(iy, i); + for (int ix = 0; ix < Qx::nrc_y; ++ix) acc[Qx::nrc_y*iy + ix] = Qx::Base::acc(acc[Qx::nrc_y*iy + ix], yv, xv[ix]); + } + } + } + for (int iy = 0; iy < Qy::nrc_y; ++iy) for (int ix = 0; ix < Qx::nrc_y; ++ix) info.store(ix0+ix, iy, Qx::Base::hsum(acc[Qx::nrc_y*iy+ix])); +} + +template +IQK_NOINLINE void mul_mat_f16_f16_NxN(int n, const char * cx, size_t bx, int ix0, const DataInfo& info) { + assert(n%QF16Base::k_step == 0); + int nb = n/QF16Base::k_step; + QF16 y(info); + QF16 x(cx + ix0*bx, bx); + QF16Base::Data xv[nrc_x]; + QF16Base::Acc acc[nrc_x*nrc_y]; + auto yv = y.load1(0, 0); + for (int ix = 0; ix < nrc_x; ++ix) { + xv[ix] = x.load1(ix, 0); + acc[ix] = QF16Base::acc_first(yv, xv[ix]); + } + for (int iy = 1; iy < nrc_y; ++iy) { + yv = y.load1(iy, 0); + for (int ix = 0; ix < nrc_x; ++ix) acc[nrc_x*iy + ix] = QF16Base::acc_first(yv, xv[ix]); + } + for (int i = 1; i < nb; ++i) { + yv = y.load1(0, i); + for (int ix = 0; ix < nrc_x; ++ix) { + xv[ix] = x.load1(ix, i); + acc[ix] = QF16Base::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < nrc_y; ++iy) { + yv = y.load1(iy, i); + for (int ix = 0; ix < nrc_x; ++ix) acc[nrc_x*iy + ix] = QF16Base::acc(acc[nrc_x*iy + ix], yv, xv[ix]); + } + } + if constexpr (!is_multiple_of_k_step) { + int nb4 = n/4; + for (int i = (QF16Base::k_step/4)*nb; i < nb4; ++i) { + yv = y.load_tail(0, i); + for (int ix = 0; ix < nrc_x; ++ix) { + xv[ix] = x.load_tail(ix, i); + acc[ix] = QF16Base::acc(acc[ix], yv, xv[ix]); + } + for (int iy = 1; iy < nrc_y; ++iy) { + yv = y.load_tail(iy, i); + for (int ix = 0; ix < nrc_x; ++ix) acc[nrc_x*iy + ix] = QF16Base::acc(acc[nrc_x*iy + ix], yv, xv[ix]); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) for (int ix = 0; ix < nrc_x; ++ix) info.store(ix0+ix, iy, QF16Base::hsum(acc[nrc_x*iy+ix])); +} + +template typename Qx> +void mul_mat_Qx_Qy_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(n%4 == 0); + constexpr int k_nx = 5; + const char * cx = (const char *)vx; + if (n%Qx::Base::k_step == 0) { + for (int ix = 0; ix < nrc_x/k_nx; ++ix) { + mul_mat_Qx_Qy_NxN, true>(n, cx, bx, ix*k_nx, info); + } + int last_x = k_nx*(nrc_x/k_nx); + if (last_x == nrc_x) return; + int nx = nrc_x - last_x; + switch (nx) { + case 1: mul_mat_Qx_Qy_NxN, true>(n, cx, bx, last_x, info); break; + case 2: mul_mat_Qx_Qy_NxN, true>(n, cx, bx, last_x, info); break; + case 3: mul_mat_Qx_Qy_NxN, true>(n, cx, bx, last_x, info); break; + case 4: mul_mat_Qx_Qy_NxN, true>(n, cx, bx, last_x, info); break; + } + } else { + for (int ix = 0; ix < nrc_x/k_nx; ++ix) { + mul_mat_Qx_Qy_NxN, false>(n, cx, bx, ix*k_nx, info); + } + int last_x = k_nx*(nrc_x/k_nx); + if (last_x == nrc_x) return; + int nx = nrc_x - last_x; + switch (nx) { + case 1: mul_mat_Qx_Qy_NxN, false>(n, cx, bx, last_x, info); break; + case 2: mul_mat_Qx_Qy_NxN, false>(n, cx, bx, last_x, info); break; + case 3: mul_mat_Qx_Qy_NxN, false>(n, cx, bx, last_x, info); break; + case 4: mul_mat_Qx_Qy_NxN, false>(n, cx, bx, last_x, info); break; + } + } +} + +template +void mul_mat_f16_f16_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(n%4 == 0); + constexpr int k_nx = 5; + const char * cx = (const char *)vx; + if (n%QF16Base::k_step == 0) { + for (int ix = 0; ix < nrc_x/k_nx; ++ix) { + mul_mat_f16_f16_NxN(n, cx, bx, ix*k_nx, info); + } + int last_x = k_nx*(nrc_x/k_nx); + if (last_x == nrc_x) return; + int nx = nrc_x - last_x; + switch (nx) { + case 1: mul_mat_f16_f16_NxN(n, cx, bx, last_x, info); break; + case 2: mul_mat_f16_f16_NxN(n, cx, bx, last_x, info); break; + case 3: mul_mat_f16_f16_NxN(n, cx, bx, last_x, info); break; + case 4: mul_mat_f16_f16_NxN(n, cx, bx, last_x, info); break; + } + } else { + for (int ix = 0; ix < nrc_x/k_nx; ++ix) { + mul_mat_f16_f16_NxN(n, cx, bx, ix*k_nx, info); + } + int last_x = k_nx*(nrc_x/k_nx); + if (last_x == nrc_x) return; + int nx = nrc_x - last_x; + switch (nx) { + case 1: mul_mat_f16_f16_NxN(n, cx, bx, last_x, info); break; + case 2: mul_mat_f16_f16_NxN(n, cx, bx, last_x, info); break; + case 3: mul_mat_f16_f16_NxN(n, cx, bx, last_x, info); break; + case 4: mul_mat_f16_f16_NxN(n, cx, bx, last_x, info); break; + } + } +} + +template +IQK_NOINLINE void mul_mat_f16_f16_Nx1(int n, const char * cx, size_t bx, int ix0, const DataInfo& info) { + assert(n%QF16Base::k_step == 0); + int nb = n/QF16Base::k_step; + QF16<1> y(info); + QF16 x(cx + ix0*bx, bx); + QF16Base::Acc acc[4*nrc_x]; + auto yv = y.loadx(0, 0); + for (int ix = 0; ix < nrc_x; ++ix) { + for (int k = 0; k < 4; ++k) { + auto xv = x.load1(ix, k); + acc[4*ix+k] = QF16Base::acc_first(yv.val[k], xv); + } + } + for (int i = 1; i < nb/4; ++i) { + yv = y.loadx(0, i); + for (int ix = 0; ix < nrc_x; ++ix) { + for (int k = 0; k < 4; ++k) { + auto xv = x.load1(ix, 4*i+k); + acc[4*ix+k] = QF16Base::acc(acc[4*ix+k], yv.val[k], xv); + } + } + } + for (int i = 4*(nb/4); i < nb; ++i) { + auto yv1 = y.load1(0, i); + for (int ix = 0; ix < nrc_x; ++ix) { + auto xv1 = x.load1(ix, i); + acc[4*ix] = QF16Base::acc(acc[4*ix], yv1, xv1); + } + } + if constexpr (!is_multiple_of_k_step) { + int nb4 = n/4; + for (int i = (QF16Base::k_step/4)*nb; i < nb4; ++i) { + auto yv1 = y.load_tail(0, i); + for (int ix = 0; ix < nrc_x; ++ix) { + auto xv1 = x.load_tail(ix, i); + acc[4*ix] = QF16Base::acc(acc[4*ix], yv1, xv1); + } + } + } + for (int ix = 0; ix < nrc_x; ++ix) { + auto v1 = vaddq_f16(acc[4*ix+0], acc[4*ix+1]); + auto v2 = vaddq_f16(acc[4*ix+2], acc[4*ix+3]); + info.store(ix0+ix, 0, QF16Base::hsum(vaddq_f16(v1, v2))); + } +} + +// At least on my M2-Max the version below, which does the multiplication row-by-row, is faster. +// But let's keep this version commented out for now. +//void mul_mat_f16_f16_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { +// GGML_ASSERT(n%4 == 0); +// constexpr int k_nx = 2; +// const char * cx = (const char *)vx; +// if (n%QF16Base::k_step == 0) { +// for (int ix = 0; ix < nrc_x/k_nx; ++ix) { +// mul_mat_f16_f16_Nx1(n, cx, bx, ix*k_nx, info); +// } +// int last_x = k_nx*(nrc_x/k_nx); +// if (last_x == nrc_x) return; +// int nx = nrc_x - last_x; +// switch (nx) { +// case 1: mul_mat_f16_f16_Nx1<1, true>(n, cx, bx, last_x, info); break; +// //case 2: mul_mat_f16_f16_Nx1<2, true>(n, cx, bx, last_x, info); break; +// //case 3: mul_mat_f16_f16_Nx1<3, true>(n, cx, bx, last_x, info); break; +// } +// } else { +// for (int ix = 0; ix < nrc_x/k_nx; ++ix) { +// mul_mat_f16_f16_Nx1(n, cx, bx, ix*k_nx, info); +// } +// int last_x = k_nx*(nrc_x/k_nx); +// if (last_x == nrc_x) return; +// int nx = nrc_x - last_x; +// switch (nx) { +// case 1: mul_mat_f16_f16_Nx1<1, false>(n, cx, bx, last_x, info); break; +// //case 2: mul_mat_f16_f16_Nx1<2, false>(n, cx, bx, last_x, info); break; +// //case 3: mul_mat_f16_f16_Nx1<3, false>(n, cx, bx, last_x, info); break; +// } +// } +//} + +void mul_mat_f16_f16_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(n%4 == 0); + const char * cx = (const char *)vx; + if (n%QF16Base::k_step == 0) { + for (int ix = 0; ix < nrc_x; ++ix) { + mul_mat_f16_f16_Nx1<1, true>(n, cx, bx, ix, info); + } + } else { + for (int ix = 0; ix < nrc_x; ++ix) { + mul_mat_f16_f16_Nx1<1, false>(n, cx, bx, ix, info); + } + } +} + +} + +bool iqk_set_kernels_float(int ne00, int typeA, int typeB, std::array& kernels) { + + if (ne00%4 == 0) { + + if (typeA == GGML_TYPE_F16 && typeB == GGML_TYPE_F16) { + for (auto& f : kernels) f = nullptr; + kernels[0] = mul_mat_f16_f16_1; + kernels[1] = mul_mat_f16_f16_T<2>; + kernels[2] = mul_mat_f16_f16_T<3>; + kernels[3] = mul_mat_f16_f16_T<4>; + kernels[4] = mul_mat_f16_f16_T<5>; + return true; + } + else if (typeA == GGML_TYPE_BF16 && typeB == GGML_TYPE_F32) { + for (auto& f : kernels) f = nullptr; + kernels[0] = mul_mat_Qx_Qy_T, QBF16>; + kernels[1] = mul_mat_Qx_Qy_T, QBF16>; + kernels[2] = mul_mat_Qx_Qy_T, QBF16>; + kernels[3] = mul_mat_Qx_Qy_T, QBF16>; + kernels[4] = mul_mat_Qx_Qy_T, QBF16>; + return true; + } + + } + + return false; + +} + +namespace { +template +inline void mm_helper(int D, int nq, const char * cx, size_t bx, DataInfo& info, int k_step) { + constexpr int nrc_k = 6; + int krem = k_step - nrc_k*(k_step/nrc_k); + for (int iq = 0; iq < nq/nrc_q; ++iq) { + for (int ik = 0; ik < k_step/nrc_k; ++ik) { + mul_mat_f16_f16_NxN(D, cx, bx, ik*nrc_k, info); + } + if (krem > 0) { + switch (krem) { + case 1: mul_mat_f16_f16_NxN(D, cx, bx, k_step - krem, info); break; + case 2: mul_mat_f16_f16_NxN(D, cx, bx, k_step - krem, info); break; + case 3: mul_mat_f16_f16_NxN(D, cx, bx, k_step - krem, info); break; + case 4: mul_mat_f16_f16_NxN(D, cx, bx, k_step - krem, info); break; + default: mul_mat_f16_f16_NxN(D, cx, bx, k_step - krem, info); break; + } + } + info.cur_y += nrc_q; + } +} +} + +void iqk_gemm_default_floats(int D, int nq, const char * cx, size_t bx, DataInfo& info, int k_step) { + constexpr int nrc_q = 4; + mm_helper(D, nq, cx, bx, info, k_step); + if (int qrem = nq - nrc_q*(nq/nrc_q); qrem > 0) { + switch (qrem) { + case 1: mm_helper<1>(D, nq, cx, bx, info, k_step); + case 2: mm_helper<2>(D, nq, cx, bx, info, k_step); + default: mm_helper<3>(D, nq, cx, bx, info, k_step); + } + } +} + +#endif + +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_gemm_floats.h b/llama.cpp/ggml/src/iqk/iqk_gemm_floats.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_gemm_floats.h @@ -0,0 +1,13 @@ +#pragma once + +#include "iqk_common.h" + +#ifdef IQK_IMPLEMENT + +#include + +bool iqk_set_kernels_float(int ne00, int typeA, int typeB, std::array& kernels); + +void iqk_gemm_default_floats(int D, int nq, const char * vx, size_t bx, DataInfo& info, int k_step); + +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_gemm_kquants.cpp b/llama.cpp/ggml/src/iqk/iqk_gemm_kquants.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_gemm_kquants.cpp @@ -0,0 +1,4829 @@ +#include "iqk_gemm_kquants.h" +#include + +#ifdef IQK_IMPLEMENT + +#include "ggml-impl.h" + +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" +#include "ggml-quants.h" + +#ifdef __x86_64__ + +namespace { + +// Handles q4_K and q5_K scales/mins +struct Scales8K { + template + inline __m256i process_mins_and_scales(const uint8_t * data, float c, int i, const Q8& q8, __m256 * accd) { + make_q4_scales(data, utmp); + const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0])); + const __m128i mins128 = _mm256_extracti128_si256(mins_and_scales, 1); + accum_mins(mins128, q8, i, c, accd); + const __m128i sc128 = _mm256_extracti128_si256(mins_and_scales, 0); + return MM256_SET_M128I(sc128, sc128); + } +#ifdef HAVE_FANCY_SIMD + template + inline __m512i process_mins_and_scales_64(const uint8_t * data, float c, int i, const Q8& q8, __m256 * accd) { + auto scales = process_mins_and_scales(data, c, i, q8, accd); + return _mm512_inserti32x8(_mm512_castsi256_si512(scales), scales, 1); + } +#endif + template + inline void accum_mins(const __m128i& mins128, const Q8& q8, int i, float c, __m256 * accd) const { + base.accum_mins(mins128, q8, i, c, accd); + } +#ifdef HAVE_FANCY_SIMD + const __m512i shuffles512[2] = { + _mm512_set_epi64(0x0706070607060706, 0x0302030203020302, 0x0706070607060706, 0x0302030203020302, + 0x0504050405040504, 0x0100010001000100, 0x0504050405040504, 0x0100010001000100), + _mm512_set_epi64(0x0f0e0f0e0f0e0f0e, 0x0b0a0b0a0b0a0b0a, 0x0f0e0f0e0f0e0f0e, 0x0b0a0b0a0b0a0b0a, + 0x0d0c0d0c0d0c0d0c, 0x0908090809080908, 0x0d0c0d0c0d0c0d0c, 0x0908090809080908) + }; +#endif + Scales8KBase base; + + uint32_t utmp[4]; +}; + +template +inline void process_mins_16(const __m256i& all_scales, const Q8& q8, int i, float d, __m256 * accm) { + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + const __m256i prod = _mm256_madd_epi16(all_scales, q8.load_bsums(iy, i)); + accm[iy] = _mm256_fmadd_ps(_mm256_set1_ps(d * q8.scale(iy, i)), _mm256_cvtepi32_ps(prod), accm[iy]); + } +} +inline void prepare_scales_16(const __m256i& all_scales, __m256i * scales) { + const __m128i l_scales = _mm256_extracti128_si256(all_scales, 0); + const __m128i h_scales = _mm256_extracti128_si256(all_scales, 1); + scales[0] = MM256_SET_M128I(l_scales, l_scales); + scales[1] = MM256_SET_M128I(h_scales, h_scales); +} + +// Handles q3_K scales +struct ScaleQ3 { + inline __m128i make_scales(const uint16_t * s8) const { + const uint16_t * scales16 = (const uint16_t *)s8; + uint32_t aux0 = scales16[0] | (scales16[1] << 16); + uint32_t aux1 = scales16[2] | (scales16[3] << 16); + uint32_t aux2 = scales16[4] | (scales16[5] << 16); + __m128i scales128 = _mm_set_epi32( + ((aux1 >> 4) & 0x0f0f0f0f) | ((aux2 >> 2) & 0x30303030), + ((aux0 >> 4) & 0x0f0f0f0f) | ((aux2 >> 0) & 0x30303030), + (aux1 & 0x0f0f0f0f) | ((aux2 << 2) & 0x30303030), + (aux0 & 0x0f0f0f0f) | ((aux2 << 4) & 0x30303030)); + return _mm_add_epi8(scales128, m32); + } + const __m128i m32 = _mm_set1_epi8(-32); +}; + +struct Scale16 { + inline void make_scales(const __m128i& scales8, __m512i * scales) const { + auto all_scales8 = MM256_SET_M128I(scales8, scales8); + auto scales1 = _mm256_shuffle_epi8(all_scales8, shuffle1); + auto scales2 = _mm256_shuffle_epi8(all_scales8, shuffle2); + scales[0] = _mm512_cvtepi8_epi16(scales1); + scales[1] = _mm512_cvtepi8_epi16(scales2); + } + template + inline void process_mins_and_scales(int i, float c, const __m128i& mins8, const __m128i& scales8, + const Q8& q8, __m256 * accm, __m512i * scales) const { + process_mins_16(_mm256_cvtepi8_epi16(mins8), q8, i, c, accm); + make_scales(scales8, scales); + } + const __m256i shuffle1 = _mm256_set_epi32(0x07070707, 0x03030303, 0x06060606, 0x02020202, + 0x05050505, 0x01010101, 0x04040404, 0x00000000); + const __m256i shuffle2 = _mm256_set_epi32(0x0f0f0f0f, 0x0b0b0b0b, 0x0e0e0e0e, 0x0a0a0a0a, + 0x0d0d0d0d, 0x09090909, 0x0c0c0c0c, 0x08080808); +}; + +template +inline void process_mins_and_scales_16(const __m128i& scales128, const Q8& q8, int i, float d, + __m256 * accm, __m256i * scales) { + const __m256i all_scales = _mm256_cvtepi8_epi16(scales128); + process_mins_16(all_scales, q8, i, d, accm); + prepare_scales_16(all_scales, scales); +} + +inline __m256i get_scale_shuffle_8(int i) { + return _mm256_set1_epi16((2*i) | ((2*i+1) << 8)); +} + +inline void set_scales_8(const __m256i& all_scales, int j, __m256i * scales) { + scales[0] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+0)); + scales[1] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+1)); + scales[2] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+2)); + scales[3] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_8(4*j+3)); +} + +inline __m256i get_scale_shuffle_16(int i) { + static const uint8_t k_shuffle[128] = { + 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, + 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 4, 5, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 7, + 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 8, 9, 10,11,10,11,10,11,10,11,10,11,10,11,10,11,10,11, + 12,13,12,13,12,13,12,13,12,13,12,13,12,13,12,13, 14,15,14,15,14,15,14,15,14,15,14,15,14,15,14,15, + }; + return _mm256_loadu_si256((const __m256i*)k_shuffle + i); +} + +inline void set_scales_16(const __m256i& all_scales, __m256i * scales) { + scales[0] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(0)); + scales[1] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(1)); + scales[2] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(2)); + scales[3] = _mm256_shuffle_epi8(all_scales, get_scale_shuffle_16(3)); +} + +struct ScaleIQ4XS { + inline __m128i make_scales(const uint32_t scales_l, const uint16_t scales_h) { + uint32_t tmp32 = scales_h | (scales_h << 14); + const __m128i sh = _mm_slli_epi16(_mm_and_si128(_mm_srlv_epi32(_mm_set1_epi32(tmp32), hshift), hmask), 4); + const __m128i sl = _mm_and_si128(_mm_srlv_epi32(_mm_set1_epi32(scales_l), lshift), lmask); + return _mm_add_epi16(_mm_or_si128(sh, _mm_cvtepi8_epi16(_mm_shuffle_epi8(sl, lshuffle))), m32); + } + const __m128i hshift = _mm_set_epi32(12, 8, 4, 0); + const __m128i lshift = _mm_set_epi32(4, 0, 4, 0); + const __m128i hmask = _mm_set1_epi16(0x03); + const __m128i lmask = _mm_set1_epi8(0xf); + const __m128i lshuffle = _mm_set_epi32(0x07030602, 0x05010400, 0x07030602, 0x05010400); + const __m128i m32 = _mm_set1_epi16(-32); +}; + +#ifdef HAVE_FANCY_SIMD +//====================================== Zen4 ================================================== + +struct HighBit5 { + inline void apply(const uint8_t * h, Q4Bits& bits) { + auto hbits256 = _mm256_loadu_si256((const __m256i *)h); + auto hbits = _mm512_inserti32x8(_mm512_castsi256_si512(hbits256), _mm256_srli_epi16(hbits256, 1), 1); + bits.values[0] = _mm512_or_si512(bits.values[0], _mm512_and_si512(_mm512_slli_epi16(hbits, 4), mh)); + bits.values[1] = _mm512_or_si512(bits.values[1], _mm512_and_si512(_mm512_slli_epi16(hbits, 2), mh)); + bits.values[2] = _mm512_or_si512(bits.values[2], _mm512_and_si512(hbits, mh)); + bits.values[3] = _mm512_or_si512(bits.values[3], _mm512_and_si512(_mm512_srli_epi16(hbits, 2), mh)); + } + const __m512i mh = _mm512_set1_epi8(0x10); +}; + +struct HighBit3 { + inline void apply(const uint8_t * h, Q2Bits& bits) { + auto hbits256 = _mm256_loadu_si256((const __m256i *)h); + auto hbits = _mm512_inserti32x8(_mm512_castsi256_si512(hbits256), _mm256_srli_epi16(hbits256, 1), 1); + bits.values[0] = _mm512_or_si512(bits.values[0], _mm512_and_si512(_mm512_slli_epi16(hbits, 2), mh)); + bits.values[1] = _mm512_or_si512(bits.values[1], _mm512_and_si512(hbits, mh)); + bits.values[2] = _mm512_or_si512(bits.values[2], _mm512_and_si512(_mm512_srli_epi16(hbits, 2), mh)); + bits.values[3] = _mm512_or_si512(bits.values[3], _mm512_and_si512(_mm512_srli_epi16(hbits, 4), mh)); + } + const __m512i mh = _mm512_set1_epi8(0x04); +}; + + +template +inline void compute_block(int iy, int i, float d, const Q8& q8, const __m512i * values, const __m512i * scales, __m512 * accd) { + const __m512i p1 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[0], q8.load_quants64(iy, i, 0)); + const __m512i p2 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[1], q8.load_quants64(iy, i, 1)); + const __m512i p3 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[2], q8.load_quants64(iy, i, 2)); + const __m512i p4 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[3], q8.load_quants64(iy, i, 3)); + auto sumi = _mm512_dpwssd_epi32(_mm512_setzero_si512(), scales[0], _mm512_packs_epi32(p1, p2)); + sumi = _mm512_dpwssd_epi32(sumi, scales[1], _mm512_packs_epi32(p3, p4)); + accd[iy] = _mm512_fmadd_ps(_mm512_set1_ps(d*q8.scale(iy, i)), _mm512_cvtepi32_ps(sumi), accd[iy]); +} + +struct DequantizerQ2K final : public BaseDequantizer { + DequantizerQ2K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + template + inline void new_block(int i, const Q8& q8, __m256 * accm, __m512i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + bits.prepare(x[i].qs); + const __m128i mins_and_scales = _mm_loadu_si128((const __m128i*)x[i].scales); + const __m128i scales8 = _mm_and_si128(mins_and_scales, m4); + const __m128i mins8 = _mm_and_si128(_mm_srli_epi16(mins_and_scales, 4), m4); + sc16.process_mins_and_scales(i, -GGML_FP16_TO_FP32(x[i].dmin), mins8, scales8, q8, accm, scales); + } + + Q2Bits bits; + Scale16 sc16; + const __m128i m4 = _mm_set1_epi8(0xf); + +}; + +struct DequantizerQ3K final : public BaseDequantizer { + DequantizerQ3K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + template + inline void new_block(int i, const Q8& q8, __m256 * accm, __m512i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + bits.prepare(x[i].qs); + hbits.apply(x[i].hmask, bits); + auto scales128 = sc3.make_scales((const uint16_t *)x[i].scales); + sc16.process_mins_and_scales(i, -4.f*d, scales128, scales128, q8, accm, scales); + } + + Q2Bits bits; + HighBit3 hbits; + ScaleQ3 sc3; + Scale16 sc16; + const __m128i m4 = _mm_set1_epi8(0xf); + const __m128i m32 = _mm_set1_epi8(-32); +}; + +struct DequantizerQ4K final : public BaseDequantizer { + DequantizerQ4K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + template + inline void new_block(int i, const Q8& q8, __m256 * accd, __m512i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + bits.prepare(x[i].qs); + auto all_scales = s8k.process_mins_and_scales_64(x[i].scales, -GGML_FP16_TO_FP32(x[i].dmin), i, q8, accd); + scales[0] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[0]); + scales[1] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[1]); + } + + Q4Bits bits; + Scales8K s8k; +}; + +struct DequantizerQ5K final : public BaseDequantizer { + DequantizerQ5K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + template + inline void new_block(int i, const Q8& q8, __m256 * accd, __m512i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + bits.prepare(x[i].qs); + hbits.apply(x[i].qh, bits); + auto all_scales = s8k.process_mins_and_scales_64(x[i].scales, -GGML_FP16_TO_FP32(x[i].dmin), i, q8, accd); + scales[0] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[0]); + scales[1] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[1]); + } + + Q4Bits bits; + HighBit5 hbits; + Scales8K s8k; +}; + +struct DequantizerQ6K final : public BaseDequantizer { + DequantizerQ6K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + template + inline void new_block(int i, const Q8& q8, __m256 * accm, __m512i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + bits.prepare64(x[i].ql); + add_high_bits(x[i].qh, bits); + auto scales128 = _mm_loadu_si128((const __m128i *)x[i].scales); + sc16.process_mins_and_scales(i, -32.f*d, scales128, scales128, q8, accm, scales); + } + + inline void add_high_bits(const uint8_t * qh, Q4Bits& bits) const { + auto hbits = _mm512_loadu_si512((const __m512i *)qh); + auto tmp1 = _mm512_and_si512(_mm512_slli_epi16(hbits, 4), mh); + auto tmp2 = _mm512_and_si512(_mm512_slli_epi16(hbits, 2), mh); + bits.values[0] = _mm512_or_si512(bits.values[0], _mm512_permutex2var_epi64(tmp1, bits.perm.permute1, tmp2)); + bits.values[2] = _mm512_or_si512(bits.values[2], _mm512_permutex2var_epi64(tmp1, bits.perm.permute2, tmp2)); + tmp1 = _mm512_and_si512(hbits, mh); + tmp2 = _mm512_and_si512(_mm512_srli_epi16(hbits, 2), mh); + bits.values[1] = _mm512_or_si512(bits.values[1], _mm512_permutex2var_epi64(tmp1, bits.perm.permute1, tmp2)); + bits.values[3] = _mm512_or_si512(bits.values[3], _mm512_permutex2var_epi64(tmp1, bits.perm.permute2, tmp2)); + } + + Q4Bits bits; + HighBit3 hbits; + Scale16 sc16; + + const __m512i mh = _mm512_set1_epi8(0x30); + +}; + +struct DequantizerIQ4XS final : public BaseDequantizer { + DequantizerIQ4XS(const void * vx, size_t bx) : BaseDequantizer(vx, bx), values(load_iq4nl_values_512()) {} + template + inline void new_block(int i, const Q8& q8, __m256 * accd, __m512i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + prepare(x[i].qs); + auto scales128 = siq4.make_scales(*(const uint32_t *)x[i].scales_l, x[i].scales_h); + s8k.accum_mins(scales128, q8, i, -128.f*d, accd); + auto scales256 = MM256_SET_M128I(scales128, scales128); + auto all_scales = _mm512_inserti32x8(_mm512_castsi256_si512(scales256), scales256, 1); + scales[0] = _mm512_shuffle_epi8(all_scales, shuffles[0]); + scales[1] = _mm512_shuffle_epi8(all_scales, shuffles[1]); + scales[2] = _mm512_shuffle_epi8(all_scales, shuffles[2]); + scales[3] = _mm512_shuffle_epi8(all_scales, shuffles[3]); + } + inline void prepare(const uint8_t * q4) { + bits.prepare64(q4); + // We now have in bits.valuse[0]: 0...15, 32...47, 64...79, 96...111 + // bits.valuse[1]: 16..31, 48...63, 80...95, 112..127 + // etc. + auto tmp = _mm512_permutex2var_epi64(bits.values[0], permute1, bits.values[1]); + bits.values[1] = _mm512_shuffle_epi8(values, _mm512_permutex2var_epi64(bits.values[0], permute2, bits.values[1])); + bits.values[0] = _mm512_shuffle_epi8(values, tmp); + tmp = _mm512_permutex2var_epi64(bits.values[2], permute1, bits.values[3]); + bits.values[3] = _mm512_shuffle_epi8(values, _mm512_permutex2var_epi64(bits.values[2], permute2, bits.values[3])); + bits.values[2] = _mm512_shuffle_epi8(values, tmp); + } + + Q4Bits bits; + Scales8KBase s8k; + ScaleIQ4XS siq4; + const __m512i values; + const __m512i permute1 = _mm512_set_epi64(11, 10, 3, 2, 9, 8, 1, 0); + const __m512i permute2 = _mm512_set_epi64(15, 14, 7, 6, 13, 12, 5, 4); + const __m512i shuffles[4] = { + _mm512_inserti32x8(_mm512_set1_epi16(0x0100), _mm256_set1_epi16(0x0302), 1), + _mm512_inserti32x8(_mm512_set1_epi16(0x0504), _mm256_set1_epi16(0x0706), 1), + _mm512_inserti32x8(_mm512_set1_epi16(0x0908), _mm256_set1_epi16(0x0b0a), 1), + _mm512_inserti32x8(_mm512_set1_epi16(0x0d0c), _mm256_set1_epi16(0x0f0e), 1), + }; +}; + +template +static void mul_mat_qX_K_q8_K_AVX512_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + constexpr int k_nx = 2; + + Q8<1> q8(info); + + Dequantizer deq1(vx, bx); + Dequantizer deq2(vx, bx); + + Dequantizer * deq[k_nx]; + deq[0] = &deq1; + deq[1] = &deq2; + + __m512i scales[2*k_nx]; + + for (int ix = 0; ix < nrc_x; ++ix) { + + auto accd = _mm512_setzero_ps(); + auto accm = _mm256_setzero_ps(); + + for (int kx = 0; kx < k_nx; ++kx) deq[kx]->new_row(ix); + + for (int i = 0; i < nb/k_nx; ++i) { + + for (int kx = 0; kx < k_nx; ++kx) deq[kx]->new_block(k_nx*i+kx, q8, &accm, scales+2*kx); + + for (int kx = 0; kx < k_nx; ++kx) { + compute_block(0, k_nx*i+kx, deq[kx]->d, q8, deq[kx]->bits.values, scales+2*kx, &accd); + } + + } + if (2*(nb/2) < nb) { + int i0 = 2*(nb/2); + deq[0]->new_block(i0, q8, &accm, scales); + compute_block(0, i0, deq[0]->d, q8, deq[0]->bits.values, scales, &accd); + } + + auto sum256 = _mm256_add_ps(_mm512_castps512_ps256(accd), _mm512_extractf32x8_ps(accd, 1)); + info.store(ix, 0, hsum_float_8(_mm256_add_ps(accm, sum256))); + } +} + +template +static void mul_mat_qX_K_q8_K_AVX512(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + Q8 q8(info); + + Dequantizer deq(vx, bx); + + __m256 accm[nrc_y]; + __m512 accd[nrc_y]; + __m512i scales[2]; + + for (int ix = 0; ix < nrc_x; ++ix) { + + for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm512_setzero_ps(); + for (int iy = 0; iy < nrc_y; ++iy) accm[iy] = _mm256_setzero_ps(); + + deq.new_row(ix); + + for (int i = 0; i < nb; ++i) { + + deq.new_block(i, q8, accm, scales); + + for (int iy = 0; iy < nrc_y; ++iy) { + const __m512i p1 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[0], q8.load_quants64(iy, i, 0)); + const __m512i p2 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[1], q8.load_quants64(iy, i, 1)); + const __m512i p3 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[2], q8.load_quants64(iy, i, 2)); + const __m512i p4 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), deq.bits.values[3], q8.load_quants64(iy, i, 3)); + auto sumi = _mm512_dpwssd_epi32(_mm512_setzero_si512(), scales[0], _mm512_packs_epi32(p1, p2)); + sumi = _mm512_dpwssd_epi32(sumi, scales[1], _mm512_packs_epi32(p3, p4)); + accd[iy] = _mm512_fmadd_ps(_mm512_set1_ps(deq.d*q8.scale(iy, i)), _mm512_cvtepi32_ps(sumi), accd[iy]); + } + + } + + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum256 = _mm256_add_ps(_mm512_castps512_ps256(accd[iy]), _mm512_extractf32x8_ps(accd[iy], 1)); + info.store(ix, iy, hsum_float_8(_mm256_add_ps(accm[iy], sum256))); + } + + } +} + +template +static void mul_mat_iqX_k_q8_K_AVX512(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + Q8 q8(info); + + Dequantizer deq(vx, bx); + + __m256 accm[nrc_y]; + __m512 accd[nrc_y]; + __m512i scales[4]; + + for (int ix = 0; ix < nrc_x; ++ix) { + + for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm512_setzero_ps(); + for (int iy = 0; iy < nrc_y; ++iy) accm[iy] = _mm256_setzero_ps(); + + deq.new_row(ix); + + for (int i = 0; i < nb; ++i) { + + deq.new_block(i, q8, accm, scales); + + for (int iy = 0; iy < nrc_y; ++iy) { + const __m512i p1 = _mm512_maddubs_epi16(deq.bits.values[0], q8.load_quants64(iy, i, 0)); + const __m512i p2 = _mm512_maddubs_epi16(deq.bits.values[1], q8.load_quants64(iy, i, 1)); + const __m512i p3 = _mm512_maddubs_epi16(deq.bits.values[2], q8.load_quants64(iy, i, 2)); + const __m512i p4 = _mm512_maddubs_epi16(deq.bits.values[3], q8.load_quants64(iy, i, 3)); + auto sumi = _mm512_dpwssd_epi32(_mm512_dpwssd_epi32(_mm512_dpwssd_epi32(_mm512_dpwssd_epi32(_mm512_setzero_si512(), + p1, scales[0]), p2, scales[1]), p3, scales[2]), p4, scales[3]); + accd[iy] = _mm512_fmadd_ps(_mm512_set1_ps(deq.d*q8.scale(iy, i)), _mm512_cvtepi32_ps(sumi), accd[iy]); + } + + } + + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum256 = _mm256_add_ps(_mm512_castps512_ps256(accd[iy]), _mm512_extractf32x8_ps(accd[iy], 1)); + info.store(ix, iy, hsum_float_8(_mm256_add_ps(accm[iy], sum256))); + } + + } +} + +#else +//====================================== AVX2 ================================================== + +struct HighBit5 { + inline void load(const uint8_t * h) { hbits = _mm256_loadu_si256((const __m256i *)h); } + inline void apply(Q4Bits& bits, bool do_shift) { + bits.values[0] = _mm256_or_si256(bits.values[0], _mm256_and_si256(_mm256_slli_epi16(hbits, 4), mh)); + bits.values[1] = _mm256_or_si256(bits.values[1], _mm256_and_si256(_mm256_slli_epi16(hbits, 3), mh)); + bits.values[2] = _mm256_or_si256(bits.values[2], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + bits.values[3] = _mm256_or_si256(bits.values[3], _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh)); + if (do_shift) { + hbits = _mm256_srli_epi16(hbits, 4); + } + } + const __m256i mh = _mm256_set1_epi8(0x10); + __m256i hbits; +}; + +struct HighBit3 { + inline void load(const uint8_t * h) { hbits = _mm256_loadu_si256((const __m256i *)h); } + inline void apply(Q2Bits& bits, bool do_shift) { + bits.values[0] = _mm256_or_si256(bits.values[0], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + bits.values[1] = _mm256_or_si256(bits.values[1], _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh)); + bits.values[2] = _mm256_or_si256(bits.values[2], _mm256_and_si256(hbits, mh)); + bits.values[3] = _mm256_or_si256(bits.values[3], _mm256_and_si256(_mm256_srli_epi16(hbits, 1), mh)); + if (do_shift) { + hbits = _mm256_srli_epi16(hbits, 4); + } + } + const __m256i mh = _mm256_set1_epi8(0x04); + __m256i hbits; +}; + +template +inline void compute_block(int iy, int i, float d, const Q8& q8, const __m512i * values, const __m512i * scales, __m512 * accd) { + const __m512i p1 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[0], q8.load_quants64(iy, i, 0)); + const __m512i p2 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[1], q8.load_quants64(iy, i, 1)); + const __m512i p3 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[2], q8.load_quants64(iy, i, 2)); + const __m512i p4 = _mm512_dpbusd_epi32(_mm512_setzero_si512(), values[3], q8.load_quants64(iy, i, 3)); + auto sumi = _mm512_dpwssd_epi32(_mm512_setzero_si512(), scales[0], _mm512_packs_epi32(p1, p2)); + sumi = _mm512_dpwssd_epi32(sumi, scales[1], _mm512_packs_epi32(p3, p4)); + accd[iy] = _mm512_fmadd_ps(_mm512_set1_ps(d*q8.scale(iy, i)), _mm512_cvtepi32_ps(sumi), accd[iy]); +} + +struct DequantizerQ2K final : public BaseDequantizer { + DequantizerQ2K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + + template + inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + const __m128i mins_and_scales = _mm_loadu_si128((const __m128i*)x[i].scales); + const __m128i scales8 = _mm_and_si128(mins_and_scales, m4); + const __m128i mins8 = _mm_and_si128(_mm_srli_epi16(mins_and_scales, 4), m4); + process_mins_16(_mm256_cvtepi8_epi16(mins8), q8, i, -GGML_FP16_TO_FP32(x[i].dmin), accm); + prepare_scales_16(_mm256_cvtepi8_epi16(scales8), scales); + } + inline void prepare(int i, int j) { + bits.prepare(x[i].qs, j); + } + + Q2Bits bits; + + const __m128i m4 = _mm_set1_epi8(0xf); +}; + +struct DequantizerQ3K final : public BaseDequantizer { + DequantizerQ3K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + + template + inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + hbits.load(x[i].hmask); + process_mins_and_scales_16(sc3.make_scales((const uint16_t *)x[i].scales), q8, i, -4.f*d, accm, scales); + } + inline void prepare(int i, int j) { + bits.prepare(x[i].qs, j); + hbits.apply(bits, j == 0); + } + + Q2Bits bits; + HighBit3 hbits; + ScaleQ3 sc3; + + const __m128i m32 = _mm_set1_epi8(-32); +}; + +struct DequantizerQ4K final : public BaseDequantizer { + DequantizerQ4K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + template + inline __m256i new_block(int i, const Q8& q8, __m256 * accd) { + d = GGML_FP16_TO_FP32(x[i].d); + return s8k.process_mins_and_scales(x[i].scales, -GGML_FP16_TO_FP32(x[i].dmin), i, q8, accd); + } + inline void prepare(int i, int j) { + bits.prepare(x[i].qs, j); + } + + Q4Bits bits; + Scales8K s8k; +}; + +struct DequantizerQ5K final : public BaseDequantizer { + DequantizerQ5K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + template + inline __m256i new_block(int i, const Q8& q8, __m256 * accd) { + d = GGML_FP16_TO_FP32(x[i].d); + hbits.load(x[i].qh); + return s8k.process_mins_and_scales(x[i].scales, -GGML_FP16_TO_FP32(x[i].dmin), i, q8, accd); + } + inline void prepare(int i, int j) { + bits.prepare(x[i].qs, j); + hbits.apply(bits, j == 0); + } + + Q4Bits bits; + HighBit5 hbits; + Scales8K s8k; +}; + +struct DequantizerQ6K final : public BaseDequantizer { + DequantizerQ6K(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + template + inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) { + d = GGML_FP16_TO_FP32(x[i].d); + process_mins_and_scales_16(_mm_loadu_si128((const __m128i *)x[i].scales), q8, i, -32.f*d, accm, scales); + } + inline void prepare(int i, int j) { + bits.prepare64(x[i].ql, j); + auto hbits = _mm256_loadu_si256((const __m256i *)x[i].qh + j); + bits.values[0] = _mm256_or_si256(bits.values[0], _mm256_and_si256(_mm256_slli_epi16(hbits, 4), mh)); + bits.values[1] = _mm256_or_si256(bits.values[1], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + bits.values[2] = _mm256_or_si256(bits.values[2], _mm256_and_si256(hbits, mh)); + bits.values[3] = _mm256_or_si256(bits.values[3], _mm256_and_si256(_mm256_srli_epi16(hbits, 2), mh)); + } + + Q4Bits bits; + const __m256i mh = _mm256_set1_epi8(0x30); +}; + +struct DequantizerIQ4XS final : public BaseDequantizer { + DequantizerIQ4XS(const void * vx, size_t bx) : BaseDequantizer(vx, bx), values(load_iq4nl_values_256()) {} + template + inline __m256i new_block(int i, const Q8& q8, __m256 * accd) { + d = GGML_FP16_TO_FP32(x[i].d); + auto scales128 = siq4.make_scales(*(const uint32_t *)x[i].scales_l, x[i].scales_h); + s8k.accum_mins(scales128, q8, i, -128.f*d, accd); + return MM256_SET_M128I(scales128, scales128); + } + inline void prepare(int i, int j) { + bits.prepare16(x[i].qs, j); + bits.values[0] = _mm256_shuffle_epi8(values, bits.values[0]); + bits.values[1] = _mm256_shuffle_epi8(values, bits.values[1]); + bits.values[2] = _mm256_shuffle_epi8(values, bits.values[2]); + bits.values[3] = _mm256_shuffle_epi8(values, bits.values[3]); + } + + Q4Bits bits; + Scales8K s8k; + ScaleIQ4XS siq4; + const __m256i values; +}; + +template +static void mul_mat_qX_K_q8_K_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + Q8 q8(info); + + Dequantizer deq(vx, bx); + + __m256 accd[nrc_y]; + __m256i scales[4]; + + for (int ix = 0; ix < nrc_x; ++ix) { + + for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm256_setzero_ps(); + + deq.new_row(ix); + + for (int i = 0; i < nb; ++i) { + + auto all_scales = deq.new_block(i, q8, accd); + + __m256i sumi[nrc_y]; + + for (int j = 0; j < QK_K/128; ++j) { + + deq.prepare(i, j); + + set_scales_8(all_scales, j, scales); + + multiply_add(deq.bits, scales, j, i, q8, sumi); + + } + + for (int iy = 0; iy < nrc_y; ++iy) { + const __m256 vd = _mm256_set1_ps(deq.d*q8.scale(iy, i)); + accd[iy] = _mm256_fmadd_ps(vd, _mm256_cvtepi32_ps(sumi[iy]), accd[iy]); + } + + } + + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, hsum_float_8(accd[iy])); + } + + } +} + +template +static void mul_mat_qY_K_q8_K_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n%QK_K == 0); + const int nb = n/QK_K; + + Q8 q8(info); + + __m256i all_scales[2]; + __m256i scales[4]; + __m256 accd[nrc_y]; + + Dequantizer deq(vx, bx); + + for (int ix = 0; ix < nrc_x; ++ix) { + + deq.new_row(ix); + + for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm256_setzero_ps(); + + for (int i = 0; i < nb; ++i) { + + deq.new_block(i, q8, accd, all_scales); + + __m256i sumi[nrc_y]; + + for (int j = 0; j < QK_K/128; ++j) { + deq.prepare(i, j); + set_scales_16(all_scales[j], scales); + multiply_add(deq.bits, scales, j, i, q8, sumi); + } + + for (int iy = 0; iy < nrc_y; ++iy) { + accd[iy] = _mm256_fmadd_ps(_mm256_set1_ps(deq.d*q8.scale(iy, i)), _mm256_cvtepi32_ps(sumi[iy]), accd[iy]); + } + + } + + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, hsum_float_8(accd[iy])); + } + + } + +} + +#endif + +// inline __m256i process_mins_and_scales(const uint8_t * data, float c, int i, const Q8& q8, __m256 * accd) { +// make_q4_scales(data, utmp); +// const __m256i mins_and_scales = _mm256_cvtepu8_epi16(_mm_set_epi32(utmp[3], utmp[2], utmp[1], utmp[0])); +// const __m128i mins128 = _mm256_extracti128_si256(mins_and_scales, 1); +// accum_mins(mins128, q8, i, c, accd); +// const __m128i sc128 = _mm256_extracti128_si256(mins_and_scales, 0); +// return MM256_SET_M128I(sc128, sc128); +// } +// +// inline void new_block(int i, const Q8& q8, __m256 * accd, __m512i * scales) { +// d = GGML_FP16_TO_FP32(x[i].d); +// bits.prepare(x[i].qs); +// auto all_scales = s8k.process_mins_and_scales_64(x[i].scales, -GGML_FP16_TO_FP32(x[i].dmin), i, q8, accd); +// scales[0] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[0]); +// scales[1] = _mm512_shuffle_epi8(all_scales, s8k.shuffles512[1]); +// } + + +struct Q4Bits_AVX2 { + inline void prepare(const uint8_t * q4, int j) { + auto q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+0); + values[0] = _mm256_and_si256(q4bits, ml); + values[1] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); + q4bits = _mm256_loadu_si256((const __m256i*)q4 + 2*j+1); + values[2] = _mm256_and_si256(q4bits, ml); + values[3] = _mm256_and_si256(_mm256_srli_epi16(q4bits, 4), ml); + } + __m256i values[4]; + const __m256i ml = _mm256_set1_epi8(0xf); +}; + +struct DequantizerQ4K_AVX2 final : public BaseDequantizer { + DequantizerQ4K_AVX2(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + inline void prepare(int i, int j) { + bits.prepare(x[i].qs, j); + } + Q4Bits_AVX2 bits; +}; + +struct DequantizerQ5K_AVX2 final : public BaseDequantizer { + DequantizerQ5K_AVX2(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + inline void prepare(int i, int j) { + bits.prepare(x[i].qs, j); + hbits = j == 0 ? _mm256_loadu_si256((const __m256i *)x[i].qh) : _mm256_srli_epi16(hbits, 4); + apply_hbits(); + } + inline void apply_hbits() { + bits.values[0] = _mm256_or_si256(bits.values[0], _mm256_and_si256(_mm256_slli_epi16(hbits, 4), mh)); + bits.values[1] = _mm256_or_si256(bits.values[1], _mm256_and_si256(_mm256_slli_epi16(hbits, 3), mh)); + bits.values[2] = _mm256_or_si256(bits.values[2], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + bits.values[3] = _mm256_or_si256(bits.values[3], _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh)); + } + + const __m256i mh = _mm256_set1_epi8(0x10); + Q4Bits_AVX2 bits; + __m256i hbits; +}; + +template +static void mul_mat_qX_K_q8_2_X4_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + Q8 q8(info); + + Dequantizer deq(vx, bx); + + uint32_t utmp[4]; + __m256 accd[nrc_y]; + __m256 scales[2]; + float d8[8*nrc_y]; + + for (int ix = 0; ix < nrc_x; ++ix) { + + for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm256_setzero_ps(); + + deq.new_row(ix); + + for (int i = 0; i < nb; ++i) { + + deq.d = GGML_FP16_TO_FP32(deq.x[i].d); + auto vm = _mm256_cvtph_ps(_mm_set1_epi16(deq.x[i].dmin)); + make_q4_scales(deq.x[i].scales, utmp); + auto mins = _mm256_mul_ps(vm, _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *)(utmp + 2))))); + mins = _mm256_mul_ps(_mm256_set1_ps(-1.f), mins); + for (int iy = 0; iy < nrc_y; ++iy) { + auto d4_1 = _mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][2*i+0].d))); + auto d4_2 = _mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][2*i+1].d))); + auto dy = _mm256_castsi256_ps(_mm256_slli_epi32(MM256_SET_M128I(d4_2, d4_1), 16)); + _mm256_storeu_ps(d8 + 8*iy, dy); + auto m4_1 = _mm_cvtepi16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][2*i+0].d+4))); + auto m4_2 = _mm_cvtepi16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][2*i+1].d+4))); + auto myi = MM256_SET_M128I(m4_2, m4_1); + auto my = _mm256_mul_ps(dy, _mm256_cvtepi32_ps(myi)); + accd[iy] = _mm256_fmadd_ps(my, mins, accd[iy]); + } + + auto all_scales = _mm256_mul_ps(_mm256_set1_ps(deq.d), _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *)utmp)))); + scales[0] = _mm256_set_m128(_mm256_castps256_ps128(all_scales), _mm256_castps256_ps128(all_scales)); + auto scales_h = _mm256_extractf128_ps(all_scales, 1); + scales[1] = _mm256_set_m128(scales_h, scales_h); + + for (int j = 0; j < QK_K/128; ++j) { + + deq.prepare(i, j); + + for (int iy = 0; iy < nrc_y; ++iy) { + const block_q8_2_x4& y = q8.y[iy][2*i+j]; +#ifdef HAVE_FANCY_SIMD + auto sumi1 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), deq.bits.values[0], _mm256_loadu_si256((const __m256i*)y.qs+0)); + auto sumi2 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), deq.bits.values[1], _mm256_loadu_si256((const __m256i*)y.qs+1)); + auto sumi3 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), deq.bits.values[2], _mm256_loadu_si256((const __m256i*)y.qs+2)); + auto sumi4 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), deq.bits.values[3], _mm256_loadu_si256((const __m256i*)y.qs+3)); + sumi1 = _mm256_add_epi32(_mm256_unpacklo_epi32(sumi1, sumi2), _mm256_unpackhi_epi32(sumi1, sumi2)); + sumi3 = _mm256_add_epi32(_mm256_unpacklo_epi32(sumi3, sumi4), _mm256_unpackhi_epi32(sumi3, sumi4)); + sumi1 = _mm256_add_epi32(_mm256_unpacklo_epi64(sumi1, sumi3), _mm256_unpackhi_epi64(sumi1, sumi3)); +#else + auto sumi1 = _mm256_maddubs_epi16(deq.bits.values[0], _mm256_loadu_si256((const __m256i*)y.qs+0)); + auto sumi2 = _mm256_maddubs_epi16(deq.bits.values[1], _mm256_loadu_si256((const __m256i*)y.qs+1)); + auto sumi3 = _mm256_maddubs_epi16(deq.bits.values[2], _mm256_loadu_si256((const __m256i*)y.qs+2)); + auto sumi4 = _mm256_maddubs_epi16(deq.bits.values[3], _mm256_loadu_si256((const __m256i*)y.qs+3)); + sumi1 = _mm256_add_epi16(_mm256_unpacklo_epi32(sumi1, sumi2), _mm256_unpackhi_epi32(sumi1, sumi2)); + sumi3 = _mm256_add_epi16(_mm256_unpacklo_epi32(sumi3, sumi4), _mm256_unpackhi_epi32(sumi3, sumi4)); + sumi1 = _mm256_add_epi16(_mm256_unpacklo_epi64(sumi1, sumi3), _mm256_unpackhi_epi64(sumi1, sumi3)); + sumi1 = _mm256_madd_epi16(_mm256_set1_epi16(1), sumi1); +#endif + auto dy4 = _mm_loadu_ps(d8 + 8*iy + 4*j); + auto d4d8 = _mm256_mul_ps(scales[j], _mm256_set_m128(dy4, dy4)); + accd[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi1), accd[iy]); + } + + } + + } + + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, hsum_float_8(accd[iy])); + } + + } +} + +struct DequantizerQ6K_AVX2 final : public BaseDequantizer { + DequantizerQ6K_AVX2(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + inline void prepare(int i, int j) { + auto lbits1 = _mm256_loadu_si256((const __m256i *)x[i].ql + 2*j+0); + auto lbits2 = _mm256_loadu_si256((const __m256i *)x[i].ql + 2*j+1); + auto hbits = _mm256_loadu_si256((const __m256i *)x[i].qh + j); + bits.values[0] = _mm256_or_si256(_mm256_and_si256(lbits1, bits.ml), _mm256_and_si256(_mm256_slli_epi16(hbits, 4), mh)); + bits.values[1] = _mm256_or_si256(_mm256_and_si256(lbits2, bits.ml), _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + bits.values[2] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits1, 4), bits.ml), _mm256_and_si256(hbits, mh)); + bits.values[3] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits2, 4), bits.ml), _mm256_and_si256(_mm256_srli_epi16(hbits, 2), mh)); + } + inline void prepare_signed(int i, int j, __m256i * us) { + prepare(i, j); + for (int k = 0; k < 4; ++k) { + bits.values[k] = _mm256_add_epi8(bits.values[k], _mm256_set1_epi8(-32)); + us[k] = _mm256_sign_epi8(bits.values[k], bits.values[k]); + } + } + inline __m256i make_scales(int i) const { + return _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i *)x[i].scales)); + } + + const __m256i mh = _mm256_set1_epi8(0x30); + Q4Bits_AVX2 bits; +}; + +struct SimpleBits { + __m256i values[4]; +}; + +struct DequantizerQ3K_AVX2 final : public BaseDequantizer { + DequantizerQ3K_AVX2(const void * vx, size_t bx) : BaseDequantizer(vx, bx) {} + + inline void prepare(int i, int j) { + hbits = j == 0 ? _mm256_loadu_si256((const __m256i *)x[i].hmask) : _mm256_srli_epi16(hbits, 4); + auto q2bits = _mm256_loadu_si256((const __m256i *)x[i].qs + j); + bits.values[0] = _mm256_and_si256(q2bits, ml); + bits.values[1] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 2), ml); + bits.values[2] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 4), ml); + bits.values[3] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 6), ml); + bits.values[0] = _mm256_or_si256(bits.values[0], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + bits.values[1] = _mm256_or_si256(bits.values[1], _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh)); + bits.values[2] = _mm256_or_si256(bits.values[2], _mm256_and_si256(hbits, mh)); + bits.values[3] = _mm256_or_si256(bits.values[3], _mm256_and_si256(_mm256_srli_epi16(hbits, 1), mh)); + //bits.values[0] = _mm256_sub_epi8(bits.values[0], _mm256_xor_si256(mh, _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh))); + //bits.values[1] = _mm256_sub_epi8(bits.values[1], _mm256_xor_si256(mh, _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh))); + //bits.values[2] = _mm256_sub_epi8(bits.values[2], _mm256_xor_si256(mh, _mm256_and_si256(hbits, mh))); + //bits.values[3] = _mm256_sub_epi8(bits.values[3], _mm256_xor_si256(mh, _mm256_and_si256(_mm256_srli_epi16(hbits, 1), mh))); + } + inline void prepare_signed(int i, int j, __m256i * us) { + prepare(i, j); + for (int k = 0; k < 4; ++k) { + bits.values[k] = _mm256_sub_epi8(bits.values[k], mh); + us[k] = _mm256_sign_epi8(bits.values[k], bits.values[k]); + } + //for (int k = 0; k < 4; ++k) { + // us[k] = _mm256_sign_epi8(bits.values[k], bits.values[k]); + //} + } + inline __m256i make_scales(int i) const { + return _mm256_cvtepi8_epi16(sc3.make_scales((const uint16_t *)x[i].scales)); + } + + ScaleQ3 sc3; + + __m256i hbits; + SimpleBits bits; + const __m256i ml = _mm256_set1_epi8(3); + const __m256i mh = _mm256_set1_epi8(4); +}; + +template +static void mul_mat_qY_K_q8_2_X4_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + Q8 q8(info); + + Dequantizer deq(vx, bx); + + __m256 accd[nrc_y]; + __m256 scales[2]; + float d8[8*nrc_y]; + __m256i us[4]; + + uint8_t k_shuff[32] = {0, 1, 4, 5, 8, 9, 12, 13, 2, 3, 6, 7, 10, 11, 14, 15, 0, 1, 4, 5, 8, 9, 12, 13, 2, 3, 6, 7, 10, 11, 14, 15}; + auto shuff = _mm256_loadu_si256((const __m256i *)k_shuff); + + for (int ix = 0; ix < nrc_x; ++ix) { + + for (int iy = 0; iy < nrc_y; ++iy) accd[iy] = _mm256_setzero_ps(); + + deq.new_row(ix); + + for (int i = 0; i < nb; ++i) { + + deq.d = GGML_FP16_TO_FP32(deq.x[i].d); + auto vd = _mm256_set1_ps(deq.d); + auto sc16 = _mm256_shuffle_epi8(deq.make_scales(i), shuff); + scales[0] = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_castsi256_si128(sc16)))); + scales[1] = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(sc16, 1)))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto d4_1 = _mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][2*i+0].d))); + auto d4_2 = _mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][2*i+1].d))); + auto dy = _mm256_castsi256_ps(_mm256_slli_epi32(MM256_SET_M128I(d4_2, d4_1), 16)); + if constexpr (nrc_y == 1) { + auto dyh = _mm256_extractf128_ps(dy, 1); + scales[0] = _mm256_mul_ps(scales[0], _mm256_set_m128(_mm256_castps256_ps128(dy), _mm256_castps256_ps128(dy))); + scales[1] = _mm256_mul_ps(scales[1], _mm256_set_m128(dyh, dyh)); + } else { + _mm256_storeu_ps(d8 + 8*iy, dy); + } + } + + for (int j = 0; j < QK_K/128; ++j) { + + deq.prepare_signed(i, j, us); + + for (int iy = 0; iy < nrc_y; ++iy) { + auto qs = q8.y[iy][2*i+j].qs; +#ifdef HAVE_FANCY_SIMD + // 0...31 + auto sumi1 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), us[0], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i*)qs+0), deq.bits.values[0])); + // 32...63 + auto sumi2 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), us[1], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i*)qs+1), deq.bits.values[1])); + // 64...95 + auto sumi3 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), us[2], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i*)qs+2), deq.bits.values[2])); + // 96...128 + auto sumi4 = _mm256_dpbusd_epi32(_mm256_setzero_si256(), us[3], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i*)qs+3), deq.bits.values[3])); + // 0...3, 32...35, 4....7, 36...39, 16...19, 48...51, 20...23, 52...56 + + // 8..11, 40...43, 12...15, 44...47, 24...27, 56...59, 28...31, 60...63 + // b0 b2 b0 b2 b1 b3 b1 b3 + sumi1 = _mm256_add_epi32(_mm256_unpacklo_epi32(sumi1, sumi2), _mm256_unpackhi_epi32(sumi1, sumi2)); + // same as above + 64, so + // b4 b6, b4 b6 b5 b7 b5 b7 + sumi3 = _mm256_add_epi32(_mm256_unpacklo_epi32(sumi3, sumi4), _mm256_unpackhi_epi32(sumi3, sumi4)); + // b0 b2 b4 b6 b1 b3 b5 b7 + + // b0 b2 b4 b6 b1 b3 b5 b7 + sumi1 = _mm256_add_epi32(_mm256_unpacklo_epi64(sumi1, sumi3), _mm256_unpackhi_epi64(sumi1, sumi3)); +#else + auto sumi1 = _mm256_maddubs_epi16(us[0], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i*)qs+0), deq.bits.values[0])); + auto sumi2 = _mm256_maddubs_epi16(us[1], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i*)qs+1), deq.bits.values[1])); + auto sumi3 = _mm256_maddubs_epi16(us[2], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i*)qs+2), deq.bits.values[2])); + auto sumi4 = _mm256_maddubs_epi16(us[3], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i*)qs+3), deq.bits.values[3])); + sumi1 = _mm256_add_epi16(_mm256_unpacklo_epi32(sumi1, sumi2), _mm256_unpackhi_epi32(sumi1, sumi2)); + sumi3 = _mm256_add_epi16(_mm256_unpacklo_epi32(sumi3, sumi4), _mm256_unpackhi_epi32(sumi3, sumi4)); + sumi1 = _mm256_add_epi16(_mm256_unpacklo_epi64(sumi1, sumi3), _mm256_unpackhi_epi64(sumi1, sumi3)); + sumi1 = _mm256_madd_epi16(_mm256_set1_epi16(1), sumi1); +#endif + if constexpr (nrc_y > 1) { + auto dy4 = _mm_loadu_ps(d8 + 8*iy + 4*j); + auto d4d8 = _mm256_mul_ps(scales[j], _mm256_set_m128(dy4, dy4)); + accd[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi1), accd[iy]); + } else { + accd[iy] = _mm256_fmadd_ps(scales[j], _mm256_cvtepi32_ps(sumi1), accd[iy]); + } + } + + } + + } + + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, hsum_float_8(accd[iy])); + } + + } +} + +template +static void mul_mat_iq4_xs_r8_q8_k_avx2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + auto m4 = _mm256_set1_epi8(0xf); + auto m30 = _mm256_set1_epi8(0x30); + auto m32 = _mm256_set1_epi8(32); +#ifndef HAVE_FANCY_SIMD + auto s_shuffle = _mm256_set_epi64x(0x0f0e0f0e0d0c0d0c, 0x0b0a0b0a09080908, 0x0706070605040504, 0x0302030201000100); + auto values128 = _mm_loadu_si128((const __m128i *)iq4k_values); + auto values = MM256_SET_M128I(values128, values128); +#else + auto values = load_iq4nl_values_256(); +#endif + int nbl = n / QK_K; + using helper_t = union { __m256i vec[2]; uint64_t val[8]; }; + helper_t h; + __m256 acc[nrc_y] = {}; + __m256i qx[4]; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_iq4_xs_r8 * iq4 = (const block_iq4_xs_r8 *)((const char *)vx + (ix+0)*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto d4 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4[ibl].d)); + auto slbits = _mm256_loadu_si256((const __m256i *)iq4[ibl].scales_l); + auto sl1 = _mm256_and_si256(slbits, m4); + auto sl2 = _mm256_and_si256(_mm256_srli_epi16(slbits, 4), m4); + auto shbits = _mm_loadu_si128((const __m128i*)iq4[ibl].scales_h); + auto sh = MM256_SET_M128I(_mm_srli_epi16(shbits, 2), shbits); + h.vec[0] = _mm256_sub_epi8(_mm256_or_si256(sl1, _mm256_and_si256(_mm256_slli_epi16(sh, 4), m30)), m32); + h.vec[1] = _mm256_sub_epi8(_mm256_or_si256(sl2, _mm256_and_si256(sh, m30)), m32); + __m256i isum[nrc_y] = {}; + for (int ib = 0; ib < QK_K/32; ++ib) { +#ifdef HAVE_FANCY_SIMD + auto iscales = _mm256_cvtepi8_epi32(_mm_set1_epi64x(h.val[ib])); + auto scales = _mm256_mul_ps(d4, _mm256_cvtepi32_ps(iscales)); + auto scales_m = _mm256_mul_ps(scales, _mm256_set1_ps(-128.f)); + for (int iy = 0; iy < nrc_y; ++iy) { + float m8 = ((const float *)q8.y[iy][ibl].bsums)[ib]; + acc[iy] = _mm256_fmadd_ps(scales_m, _mm256_set1_ps(m8), acc[iy]); + } +#else + auto iscales = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm_set1_epi64x(h.val[ib])), s_shuffle); +#endif + auto bits1 = _mm256_loadu_si256((const __m256i *)iq4[ibl].qs+4*ib+0); + auto bits2 = _mm256_loadu_si256((const __m256i *)iq4[ibl].qs+4*ib+1); + qx[0] = _mm256_shuffle_epi8(values, _mm256_and_si256(m4, bits1)); + qx[1] = _mm256_shuffle_epi8(values, _mm256_and_si256(m4, _mm256_srli_epi16(bits1, 4))); + qx[2] = _mm256_shuffle_epi8(values, _mm256_and_si256(m4, bits2)); + qx[3] = _mm256_shuffle_epi8(values, _mm256_and_si256(m4, _mm256_srli_epi16(bits2, 4))); +#ifndef HAVE_FANCY_SIMD + auto s1 = _mm256_sign_epi8(qx[0], qx[0]); + auto s2 = _mm256_sign_epi8(qx[1], qx[1]); + auto s3 = _mm256_sign_epi8(qx[2], qx[2]); + auto s4 = _mm256_sign_epi8(qx[3], qx[3]); +#endif + for (int iy = 0; iy < nrc_y; ++iy) { + auto y128 = _mm_loadu_si128((const __m128i*)q8.y[iy][ibl].qs+2*ib+0); + auto y = MM256_SET_M128I(y128, y128); +#ifdef HAVE_FANCY_SIMD + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_mullo_epi32(iscales, sumi)); +#else + auto sumi1 = _mm256_maddubs_epi16(s1, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0])); + auto sumi2 = _mm256_maddubs_epi16(s2, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1])); + auto sumi3 = _mm256_maddubs_epi16(s3, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2])); + auto sumi4 = _mm256_maddubs_epi16(s4, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qx[3])); + auto sumi = _mm256_add_epi32(_mm256_add_epi32(_mm256_madd_epi16(iscales, sumi1), _mm256_madd_epi16(iscales, sumi2)), + _mm256_add_epi32(_mm256_madd_epi16(iscales, sumi3), _mm256_madd_epi16(iscales, sumi4))); + isum[iy] = _mm256_add_epi32(isum[iy], sumi); +#endif + } + bits1 = _mm256_loadu_si256((const __m256i *)iq4[ibl].qs+4*ib+2); + bits2 = _mm256_loadu_si256((const __m256i *)iq4[ibl].qs+4*ib+3); + qx[0] = _mm256_shuffle_epi8(values, _mm256_and_si256(m4, bits1)); + qx[1] = _mm256_shuffle_epi8(values, _mm256_and_si256(m4, _mm256_srli_epi16(bits1, 4))); + qx[2] = _mm256_shuffle_epi8(values, _mm256_and_si256(m4, bits2)); + qx[3] = _mm256_shuffle_epi8(values, _mm256_and_si256(m4, _mm256_srli_epi16(bits2, 4))); +#ifndef HAVE_FANCY_SIMD + s1 = _mm256_sign_epi8(qx[0], qx[0]); + s2 = _mm256_sign_epi8(qx[1], qx[1]); + s3 = _mm256_sign_epi8(qx[2], qx[2]); + s4 = _mm256_sign_epi8(qx[3], qx[3]); +#endif + for (int iy = 0; iy < nrc_y; ++iy) { + auto y128 = _mm_loadu_si128((const __m128i*)q8.y[iy][ibl].qs+2*ib+1); + auto y = MM256_SET_M128I(y128, y128); +#ifdef HAVE_FANCY_SIMD + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_mullo_epi32(iscales, sumi)); +#else + auto sumi1 = _mm256_maddubs_epi16(s1, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0])); + auto sumi2 = _mm256_maddubs_epi16(s2, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1])); + auto sumi3 = _mm256_maddubs_epi16(s3, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2])); + auto sumi4 = _mm256_maddubs_epi16(s4, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qx[3])); + auto sumi = _mm256_add_epi32(_mm256_add_epi32(_mm256_madd_epi16(iscales, sumi1), _mm256_madd_epi16(iscales, sumi2)), + _mm256_add_epi32(_mm256_madd_epi16(iscales, sumi3), _mm256_madd_epi16(iscales, sumi4))); + isum[iy] = _mm256_add_epi32(isum[iy], sumi); +#endif + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(isum[iy]), acc[iy]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = _mm256_setzero_ps(); + } + } +} + +#ifdef HAVE_FANCY_SIMD +template +static void mul_mat_iq4_xs_r8_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + mul_mat_iq4_xs_r8_q8_k_avx2(n, vx, bx, info, nrc_x); + return; + if constexpr (nrc_y == 1){ + mul_mat_iq4_xs_r8_q8_k_avx2<1>(n, vx, bx, info, nrc_x); + } else { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + auto m4 = _mm512_set1_epi8(0xf); + auto values = load_iq4nl_values_512(); + int nbl = n / QK_K; + using helper_t = union { __m512i vec; uint32_t val[16]; }; + helper_t h; + __m512 acc[nrc_y] = {}; + __m512i isum[nrc_y] = {}; + __m512i qx[4]; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_iq4_xs_r8 * iq4l = (const block_iq4_xs_r8 *)((const char *)vx + (ix+0)*bx); + const block_iq4_xs_r8 * iq4h = (const block_iq4_xs_r8 *)((const char *)vx + (ix+4)*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto dl = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq4l[ibl].d)); + auto dh = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq4h[ibl].d)); + auto d4 = _mm512_insertf32x8(_mm512_castps256_ps512(_mm256_set_m128(dl, dl)), _mm256_set_m128(dh, dh), 1); + auto d4x64 = _mm512_mul_ps(d4, _mm512_set1_ps(-64.f)); + auto slbits_l = _mm_loadu_si128((const __m128i *)iq4l[ibl].scales_l); + auto shbits_l = _mm_loadu_si128((const __m128i *)iq4h[ibl].scales_l); + auto sl_l = MM256_SET_M128I(_mm_srli_epi16(slbits_l, 4), slbits_l); + auto sh_l = MM256_SET_M128I(_mm_srli_epi16(shbits_l, 4), shbits_l); + auto slb = _mm512_and_si512(_mm512_inserti32x8(_mm512_castsi256_si512(sl_l), sh_l, 1), m4); + auto aux64 = (const uint64_t *)iq4l[ibl].scales_h; + auto slbits_h = _mm_set_epi64x(aux64[0] >> 2, aux64[0]); + aux64 = (const uint64_t *)iq4h[ibl].scales_h; + auto shbits_h = _mm_set_epi64x(aux64[0] >> 2, aux64[0]); + auto sl_h = MM256_SET_M128I(slbits_h, _mm_slli_epi16(slbits_h, 4)); + auto sh_h = MM256_SET_M128I(shbits_h, _mm_slli_epi16(shbits_h, 4)); + auto shb = _mm512_and_si512(_mm512_inserti32x8(_mm512_castsi256_si512(sl_h), sh_h, 1), _mm512_set1_epi8(0x30)); + h.vec = _mm512_sub_epi8(_mm512_or_si512(slb, shb), _mm512_set1_epi8(32)); + for (int ib = 0; ib < QK_K/32; ++ib) { + auto iscales = _mm512_cvtepi8_epi32(_mm_blend_epi32(_mm_set1_epi32(h.val[ib+0]), _mm_set1_epi32(h.val[ib+8]), 0x0c)); + auto scales = _mm512_cvtepi32_ps(iscales); + auto scales_m = _mm512_mul_ps(scales, d4x64); + auto bits1 = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq4l[ibl].qs+2*ib+0)), + _mm256_loadu_si256((const __m256i *)iq4h[ibl].qs+2*ib+0), 1); + auto bits2 = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq4l[ibl].qs+2*ib+1)), + _mm256_loadu_si256((const __m256i *)iq4h[ibl].qs+2*ib+1), 1); + qx[0] = _mm512_shuffle_epi8(values, _mm512_and_si512(bits1, m4)); + qx[1] = _mm512_shuffle_epi8(values, _mm512_and_si512(bits2, m4)); + qx[2] = _mm512_shuffle_epi8(values, _mm512_and_si512(_mm512_srli_epi16(bits1, 4), m4)); + qx[3] = _mm512_shuffle_epi8(values, _mm512_and_si512(_mm512_srli_epi16(bits2, 4), m4)); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y8 = _mm256_loadu_si256((const __m256i*)q8.y[iy][ibl].qs+ib); + auto y = _mm512_inserti32x8(_mm512_castsi256_si512(y8), y8, 1); + auto sumi = _mm512_setzero_si512(); + sumi = _mm512_dpbusd_epi32(sumi, qx[0], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + sumi = _mm512_dpbusd_epi32(sumi, qx[1], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + sumi = _mm512_dpbusd_epi32(sumi, qx[2], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + sumi = _mm512_dpbusd_epi32(sumi, qx[3], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + isum[iy] = _mm512_add_epi32(isum[iy], _mm512_mullo_epi32(iscales, sumi)); + float m8 = ((const float *)q8.y[iy][ibl].bsums)[ib]; + acc[iy] = _mm512_fmadd_ps(scales_m, _mm512_set1_ps(m8), acc[iy]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = _mm512_fmadd_ps(_mm512_mul_ps(d4, _mm512_set1_ps(q8.scale(iy, ibl))), _mm512_cvtepi32_ps(isum[iy]), acc[iy]); + isum[iy] = _mm512_setzero_si512(); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum1 = _mm_add_ps(_mm512_extractf32x4_ps(acc[iy], 0), _mm512_extractf32x4_ps(acc[iy], 1)); + auto sum2 = _mm_add_ps(_mm512_extractf32x4_ps(acc[iy], 2), _mm512_extractf32x4_ps(acc[iy], 3)); + info.store(ix+0, iy, sum1); + info.store(ix+4, iy, sum2); + acc[iy] = _mm512_setzero_ps(); + } + } + } +} +#else +template +static void mul_mat_iq4_xs_r8_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + mul_mat_iq4_xs_r8_q8_k_avx2(n, vx, bx, info, nrc_x); +} +#endif + +template +static void mul_mat_q2_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto mxf = _mm256_set1_epi8(0xf); + auto m03 = _mm256_set1_epi8(0x03); + static const uint8_t k_shuff[32] = {0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15, 0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15}; + auto shuff = _mm256_loadu_si256((const __m256i *)k_shuff); +#ifdef HAVE_FANCY_SIMD + __m256i isum[nrc_y] = {}; +#else + auto m1 = _mm256_set1_epi16(1); +#endif + int nbl = n / QK_K; + __m256 acc[nrc_y] = {}; + __m256i qx[4]; + int8_t scales[64]; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q2_k_r4 * iq2 = (const block_q2_k_r4 *)((const char *)vx + (ix+0)*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto dm = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq2[ibl].d)); + auto d4 = _mm256_set_m128(_mm256_castps256_ps128(dm), _mm256_castps256_ps128(dm)); + auto m4 = _mm256_set_m128(_mm256_extractf128_ps(dm, 1), _mm256_extractf128_ps(dm, 1)); + m4 = _mm256_mul_ps(m4, _mm256_set1_ps(-1.f)); + auto all_scales1 = _mm256_loadu_si256((const __m256i *)iq2[ibl].scales+0); + auto all_scales2 = _mm256_loadu_si256((const __m256i *)iq2[ibl].scales+1); + auto scales1 = _mm256_and_si256(_mm256_srli_epi16(all_scales1, 4), mxf); + auto scales2 = _mm256_and_si256(_mm256_srli_epi16(all_scales2, 4), mxf); + { + auto t1 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(scales1, 0)), shuff); // blocks 0, 1, 2, 3 for each row + auto t2 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(scales1, 1)), shuff); // blocks 4, 5, 6, 7 for each row + auto t3 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(scales2, 0)), shuff); // blocks 8, 9, 10, 11 for each row + auto t4 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(scales2, 1)), shuff); // blocks 12, 13, 14, 15 for each row + auto s1 = MM256_SET_M128I(_mm256_extracti128_si256(t3, 0), _mm256_extracti128_si256(t1, 0)); // blocks 0, 1, 8, 9 + auto s2 = MM256_SET_M128I(_mm256_extracti128_si256(t3, 1), _mm256_extracti128_si256(t1, 1)); // blocks 2, 3, 10, 11 + auto s3 = MM256_SET_M128I(_mm256_extracti128_si256(t4, 0), _mm256_extracti128_si256(t2, 0)); // blocks 4, 5, 12, 13 + auto s4 = MM256_SET_M128I(_mm256_extracti128_si256(t4, 1), _mm256_extracti128_si256(t2, 1)); // blocks 6, 7, 14, 15 + for (int iy = 0; iy < nrc_y; ++iy) { + auto bsums = q8.load_bsums(iy, ibl); + auto sumi = _mm256_setzero_si256(); +#ifdef HAVE_FANCY_SIMD + sumi = _mm256_dpwssd_epi32(sumi, s1, _mm256_shuffle_epi32(bsums, 0x00)); + sumi = _mm256_dpwssd_epi32(sumi, s2, _mm256_shuffle_epi32(bsums, 0x55)); + sumi = _mm256_dpwssd_epi32(sumi, s3, _mm256_shuffle_epi32(bsums, 0xaa)); + sumi = _mm256_dpwssd_epi32(sumi, s4, _mm256_shuffle_epi32(bsums, 0xff)); + auto d8 = _mm256_set1_ps(q8.scale(iy, ibl)); + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(m4, d8), _mm256_cvtepi32_ps(sumi), acc[iy]); +#else + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s1, _mm256_shuffle_epi32(bsums, 0x00))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s2, _mm256_shuffle_epi32(bsums, 0x55))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s3, _mm256_shuffle_epi32(bsums, 0xaa))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s4, _mm256_shuffle_epi32(bsums, 0xff))); + auto d8 = _mm256_set1_ps(q8.scale(iy, ibl)); + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(m4, d8), _mm256_cvtepi32_ps(sumi), acc[iy]); + if constexpr (nrc_y == 1) { + d4 = _mm256_mul_ps(d4, d8); + } +#endif + } + } + all_scales1 = _mm256_and_si256(all_scales1, mxf); + all_scales2 = _mm256_and_si256(all_scales2, mxf); + _mm256_storeu_si256((__m256i *)scales+0, all_scales1); + _mm256_storeu_si256((__m256i *)scales+1, all_scales2); + for (int ib = 0; ib < QK_K/32; ++ib) { + auto iscales = _mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *)(scales + 8*ib))); +#ifndef HAVE_FANCY_SIMD + auto scales = _mm256_mul_ps(d4, _mm256_cvtepi32_ps(iscales)); +#endif + auto lb = _mm256_loadu_si256((const __m256i *)iq2[ibl].qs+ib); + qx[0] = _mm256_and_si256(lb, m03); + qx[1] = _mm256_and_si256(_mm256_srli_epi16(lb, 2), m03); + qx[2] = _mm256_and_si256(_mm256_srli_epi16(lb, 4), m03); + qx[3] = _mm256_and_si256(_mm256_srli_epi16(lb, 6), m03); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = _mm256_loadu_si256((const __m256i*)q8.y[iy][ibl].qs+ib); +#ifdef HAVE_FANCY_SIMD + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_mullo_epi32(iscales, sumi)); +#else + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[0], _mm256_shuffle_epi32(y, 0x00)), + _mm256_maddubs_epi16(qx[1], _mm256_shuffle_epi32(y, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[2], _mm256_shuffle_epi32(y, 0xaa)), + _mm256_maddubs_epi16(qx[3], _mm256_shuffle_epi32(y, 0xff))); + // Quants are in 0...3, so we can add add up all of them as int16_t without overflowing + auto sumi = _mm256_madd_epi16(m1, _mm256_add_epi16(sumi1, sumi2)); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(scales, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(scales, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#endif + } + } +#ifdef HAVE_FANCY_SIMD + for (int iy = 0; iy < nrc_y; ++iy) { + auto d4y = _mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(iy, ibl))); + acc[iy] = _mm256_fmadd_ps(d4y, _mm256_cvtepi32_ps(isum[iy]), acc[iy]); + isum[iy] = _mm256_setzero_si256(); + } +#endif + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm_add_ps(_mm256_castps256_ps128(acc[iy]), _mm256_extractf128_ps(acc[iy], 1)); + acc[iy] = _mm256_setzero_ps(); + info.store(ix+0, iy, sum); + } + } +} + +template +static void mul_mat_q3_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto m4 = _mm256_set1_epi8(0xf); + auto m30 = _mm256_set1_epi8(0x30); + auto m32 = _mm256_set1_epi8(32); + auto m03 = _mm256_set1_epi8(0x03); + auto m04 = _mm256_set1_epi8(0x04); + static const uint8_t k_shuff[32] = {0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15, 0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15}; + auto shuff = _mm256_loadu_si256((const __m256i *)k_shuff); +#ifdef HAVE_FANCY_SIMD + __m256i isum[nrc_y]; +#elif !defined(HAVE_VNNI256) + auto m1 = _mm256_set1_epi16(1); +#endif + int nbl = n / QK_K; + __m256 acc[nrc_y] = {}; + __m256i qx[4]; + int8_t scales[64]; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q3_k_r4 * iq3 = (const block_q3_k_r4 *)((const char *)vx + (ix+0)*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto dl = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq3[ibl].d)); + auto d4 = _mm256_set_m128(dl, dl); +#ifndef HAVE_FANCY_SIMD + if constexpr (nrc_y == 1) { + d4 = _mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(0, ibl))); + } +#endif + auto slb = _mm256_loadu_si256((const __m256i *)iq3[ibl].scales_l); + auto shbits = _mm_loadu_si128((const __m128i *)iq3[ibl].scales_h); + auto shb = MM256_SET_M128I(_mm_srli_epi16(shbits, 2), shbits); + auto scales1 = _mm256_sub_epi8(_mm256_or_si256(_mm256_and_si256(slb, m4), _mm256_and_si256(_mm256_slli_epi16(shb, 4), m30)), m32); + auto scales2 = _mm256_sub_epi8(_mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(slb, 4), m4), _mm256_and_si256(shb, m30)), m32); + _mm256_storeu_si256((__m256i *)scales+0, scales1); + _mm256_storeu_si256((__m256i *)scales+1, scales2); + { +#ifndef HAVE_FANCY_SIMD + auto min = _mm256_mul_ps(d4, _mm256_set1_ps(-4.f)); +#endif + auto t1 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(scales1, 0)), shuff); // blocks 0, 1, 2, 3 for each row + auto t2 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(scales1, 1)), shuff); // blocks 4, 5, 6, 7 for each row + auto t3 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(scales2, 0)), shuff); // blocks 8, 9, 10, 11 for each row + auto t4 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm256_extracti128_si256(scales2, 1)), shuff); // blocks 12, 13, 14, 15 for each row + auto s1 = MM256_SET_M128I(_mm256_extracti128_si256(t3, 0), _mm256_extracti128_si256(t1, 0)); // blocks 0, 1, 8, 9 + auto s2 = MM256_SET_M128I(_mm256_extracti128_si256(t3, 1), _mm256_extracti128_si256(t1, 1)); // blocks 2, 3, 10, 11 + auto s3 = MM256_SET_M128I(_mm256_extracti128_si256(t4, 0), _mm256_extracti128_si256(t2, 0)); // blocks 4, 5, 12, 13 + auto s4 = MM256_SET_M128I(_mm256_extracti128_si256(t4, 1), _mm256_extracti128_si256(t2, 1)); // blocks 6, 7, 14, 15 +#ifdef HAVE_FANCY_SIMD + s1 = _mm256_mullo_epi16(s1, _mm256_set1_epi16(-4)); + s2 = _mm256_mullo_epi16(s2, _mm256_set1_epi16(-4)); + s3 = _mm256_mullo_epi16(s3, _mm256_set1_epi16(-4)); + s4 = _mm256_mullo_epi16(s4, _mm256_set1_epi16(-4)); +#endif + for (int iy = 0; iy < nrc_y; ++iy) { + auto bsums = q8.load_bsums(iy, ibl); + auto sumi = _mm256_setzero_si256(); +#ifdef HAVE_FANCY_SIMD + sumi = _mm256_dpwssd_epi32(sumi, s1, _mm256_shuffle_epi32(bsums, 0x00)); + sumi = _mm256_dpwssd_epi32(sumi, s2, _mm256_shuffle_epi32(bsums, 0x55)); + sumi = _mm256_dpwssd_epi32(sumi, s3, _mm256_shuffle_epi32(bsums, 0xaa)); + sumi = _mm256_dpwssd_epi32(sumi, s4, _mm256_shuffle_epi32(bsums, 0xff)); + isum[iy] = sumi; +#elif defined(HAVE_VNNI256) + sumi = ggml_mm256_dpwssd_epi32(sumi, s1, _mm256_shuffle_epi32(bsums, 0x00)); + sumi = ggml_mm256_dpwssd_epi32(sumi, s2, _mm256_shuffle_epi32(bsums, 0x55)); + sumi = ggml_mm256_dpwssd_epi32(sumi, s3, _mm256_shuffle_epi32(bsums, 0xaa)); + sumi = ggml_mm256_dpwssd_epi32(sumi, s4, _mm256_shuffle_epi32(bsums, 0xff)); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(min, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(min, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#else + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s1, _mm256_shuffle_epi32(bsums, 0x00))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s2, _mm256_shuffle_epi32(bsums, 0x55))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s3, _mm256_shuffle_epi32(bsums, 0xaa))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s4, _mm256_shuffle_epi32(bsums, 0xff))); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(min, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(min, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#endif + } + } + for (int ib = 0; ib < QK_K/32; ++ib) { + auto iscales = _mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *)(scales + 8*ib))); +#ifndef HAVE_FANCY_SIMD + auto scales = _mm256_mul_ps(d4, _mm256_cvtepi32_ps(iscales)); +#endif + auto lb = _mm256_loadu_si256((const __m256i *)iq3[ibl].qs+ib); + auto hbits = _mm_loadu_si128((const __m128i *)iq3[ibl].qh+ib); + auto hb = MM256_SET_M128I(hbits, _mm_slli_epi16(hbits, 4)); + qx[0] = _mm256_or_si256(_mm256_and_si256(lb, m03), _mm256_and_si256(m04, _mm256_srli_epi16(hb, 2))); + qx[1] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lb, 2), m03), _mm256_and_si256(m04, _mm256_srli_epi16(hb, 3))); + qx[2] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lb, 4), m03), _mm256_and_si256(m04, _mm256_srli_epi16(hb, 4))); + qx[3] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lb, 6), m03), _mm256_and_si256(m04, _mm256_srli_epi16(hb, 5))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = _mm256_loadu_si256((const __m256i*)q8.y[iy][ibl].qs+ib); +#ifdef HAVE_FANCY_SIMD + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_mullo_epi32(iscales, sumi)); +#elif defined(HAVE_VNNI256) + auto sumi = _mm256_setzero_si256(); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(scales, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(scales, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#else + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[0], _mm256_shuffle_epi32(y, 0x00)), + _mm256_maddubs_epi16(qx[1], _mm256_shuffle_epi32(y, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[2], _mm256_shuffle_epi32(y, 0xaa)), + _mm256_maddubs_epi16(qx[3], _mm256_shuffle_epi32(y, 0xff))); + // Quants are in 0...8, so we can add add up all of them as int16_t without overflowing + auto sumi = _mm256_madd_epi16(m1, _mm256_add_epi16(sumi1, sumi2)); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(scales, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(scales, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#endif + + } + } +#ifdef HAVE_FANCY_SIMD + for (int iy = 0; iy < nrc_y; ++iy) { + auto d4y = _mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(iy, ibl))); + acc[iy] = _mm256_fmadd_ps(d4y, _mm256_cvtepi32_ps(isum[iy]), acc[iy]); + } +#endif + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm_add_ps(_mm256_castps256_ps128(acc[iy]), _mm256_extractf128_ps(acc[iy], 1)); + acc[iy] = _mm256_setzero_ps(); + info.store(ix+0, iy, sum); + } + } +} + +template +inline void process_min_r4_b32(int ibl, __m256 m4, __m256i mins, const Q8& q8, __m256 * acc) { + auto mins_l = _mm256_castsi256_si128(mins); + auto mins_h = _mm256_extracti128_si256(mins, 1); + auto aux1 = _mm_unpacklo_epi32(mins_l, mins_h); + auto aux2 = _mm_unpackhi_epi32(mins_l, mins_h); + auto ic1 = _mm256_cvtepi8_epi32(aux1); + auto ic2 = _mm256_cvtepi8_epi32(_mm_shuffle_epi32(aux1, 0xee)); + auto ic3 = _mm256_cvtepi8_epi32(aux2); + auto ic4 = _mm256_cvtepi8_epi32(_mm_shuffle_epi32(aux2, 0xee)); + if constexpr (nrc_y == 1) { + auto bs = _mm256_loadu_ps((const float *)q8.y[0][ibl].bsums); + auto sumf = _mm256_mul_ps(_mm256_cvtepi32_ps(ic1), _mm256_shuffle_ps(bs, bs, 0x00)); + sumf = _mm256_fmadd_ps(_mm256_cvtepi32_ps(ic2), _mm256_shuffle_ps(bs, bs, 0x55), sumf); + sumf = _mm256_fmadd_ps(_mm256_cvtepi32_ps(ic3), _mm256_shuffle_ps(bs, bs, 0xaa), sumf); + sumf = _mm256_fmadd_ps(_mm256_cvtepi32_ps(ic4), _mm256_shuffle_ps(bs, bs, 0xff), sumf); + acc[0] = _mm256_fmadd_ps(m4, sumf, acc[0]); + } else { + auto c1 = _mm256_mul_ps(m4, _mm256_cvtepi32_ps(ic1)); + auto c2 = _mm256_mul_ps(m4, _mm256_cvtepi32_ps(ic2)); + auto c3 = _mm256_mul_ps(m4, _mm256_cvtepi32_ps(ic3)); + auto c4 = _mm256_mul_ps(m4, _mm256_cvtepi32_ps(ic4)); + for (int iy = 0; iy < nrc_y; ++iy) { + auto bs = _mm256_loadu_ps((const float *)q8.y[iy][ibl].bsums); + acc[iy] = _mm256_fmadd_ps(c1, _mm256_shuffle_ps(bs, bs, 0x00), acc[iy]); + acc[iy] = _mm256_fmadd_ps(c2, _mm256_shuffle_ps(bs, bs, 0x55), acc[iy]); + acc[iy] = _mm256_fmadd_ps(c3, _mm256_shuffle_ps(bs, bs, 0xaa), acc[iy]); + acc[iy] = _mm256_fmadd_ps(c4, _mm256_shuffle_ps(bs, bs, 0xff), acc[iy]); + } + } +} + +template +static void mul_mat_q4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto mf = _mm256_set1_epi8(0xf); + auto m3 = _mm256_set1_epi8(0x30); + int nbl = n / QK_K; + union { __m256i vec; uint32_t val[8]; } hd; + __m256 acc[nrc_y] = {}; + __m256i isum[nrc_y] = {}; + __m256i qx[4]; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q4_k_r4 * iq4 = (const block_q4_k_r4 *)((const char *)vx + (ix+0)*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto dl = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4[ibl].d)); + auto d4 = _mm256_set_m128(_mm256_castps256_ps128(dl), _mm256_castps256_ps128(dl)); + auto m4 = _mm256_mul_ps(_mm256_set1_ps(-1.0f), _mm256_set_m128(_mm256_extractf128_ps(dl, 1), _mm256_extractf128_ps(dl, 1))); + auto lbits = _mm256_loadu_si256((const __m256i *)iq4[ibl].scales_l); + auto hbits128 = _mm_loadu_si128((const __m128i *)iq4[ibl].scales_h); + auto hbits = MM256_SET_M128I(hbits128, _mm_slli_epi16(hbits128, 4)); + hd.vec = _mm256_or_si256(_mm256_and_si256(lbits, mf), _mm256_and_si256(hbits, m3)); + auto mins = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits, 4), mf), _mm256_and_si256(_mm256_srli_epi16(hbits, 2), m3)); + process_min_r4_b32(ibl, m4, mins, q8, acc); + for (int ib = 0; ib < QK_K/32; ++ib) { +#ifdef HAVE_VNNI256 + auto scales_d = _mm256_cvtepi8_epi32(_mm_set1_epi32(hd.val[ib])); +#else + auto aux = _mm_set1_epi32(hd.val[ib]); + aux = _mm_cvtepu8_epi16(_mm_unpacklo_epi8(aux, aux)); + auto scales_d = MM256_SET_M128I(aux, aux); +#endif + auto bits1 = _mm256_loadu_si256((const __m256i *)iq4[ibl].qs+2*ib+0); + auto bits2 = _mm256_loadu_si256((const __m256i *)iq4[ibl].qs+2*ib+1); + qx[0] = _mm256_and_si256(bits1, mf); + qx[1] = _mm256_and_si256(bits2, mf); + qx[2] = _mm256_and_si256(_mm256_srli_epi16(bits1, 4), mf); + qx[3] = _mm256_and_si256(_mm256_srli_epi16(bits2, 4), mf); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = _mm256_loadu_si256((const __m256i*)q8.y[iy][ibl].qs+ib); +#ifdef HAVE_VNNI256 + auto sumi = _mm256_setzero_si256(); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_mullo_epi32(scales_d, sumi)); +#else + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[0], _mm256_shuffle_epi32(y, 0x00)), + _mm256_maddubs_epi16(qx[1], _mm256_shuffle_epi32(y, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[2], _mm256_shuffle_epi32(y, 0xaa)), + _mm256_maddubs_epi16(qx[3], _mm256_shuffle_epi32(y, 0xff))); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_madd_epi16(scales_d, _mm256_add_epi16(sumi1, sumi2))); +#endif + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(isum[iy]), acc[iy]); + isum[iy] = _mm256_setzero_si256(); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm_add_ps(_mm256_castps256_ps128(acc[iy]), _mm256_extractf128_ps(acc[iy], 1)); + acc[iy] = _mm256_setzero_ps(); + info.store(ix+0, iy, sum); + } + } +} + +template +static void mul_mat_q5_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto mf = _mm256_set1_epi8(0xf); + auto m10 = _mm256_set1_epi8(0x10); + auto m30 = _mm256_set1_epi8(0x30); + int nbl = n / QK_K; + union { __m256i vec; uint32_t val[8]; } hd; + __m256 acc[nrc_y] = {}; + __m256i isum[nrc_y] = {}; + __m256i qx[4]; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q5_k_r4 * iq5 = (const block_q5_k_r4 *)((const char *)vx + (ix+0)*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto dl = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq5[ibl].d)); + auto d4 = _mm256_set_m128(_mm256_castps256_ps128(dl), _mm256_castps256_ps128(dl)); + auto m4 = _mm256_mul_ps(_mm256_set1_ps(-1.0f), _mm256_set_m128(_mm256_extractf128_ps(dl, 1), _mm256_extractf128_ps(dl, 1))); + auto lbits = _mm256_loadu_si256((const __m256i *)iq5[ibl].scales_l); + auto hbits128 = _mm_loadu_si128((const __m128i *)iq5[ibl].scales_h); + auto hbits = MM256_SET_M128I(hbits128, _mm_slli_epi16(hbits128, 4)); + hd.vec = _mm256_or_si256(_mm256_and_si256(lbits, mf), _mm256_and_si256(hbits, m30)); + auto mins = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits, 4), mf), _mm256_and_si256(_mm256_srli_epi16(hbits, 2), m30)); + process_min_r4_b32(ibl, m4, mins, q8, acc); + for (int ib = 0; ib < QK_K/32; ++ib) { +#ifdef HAVE_VNNI256 + auto scales_d = _mm256_cvtepi8_epi32(_mm_set1_epi32(hd.val[ib])); +#else + auto aux = _mm_set1_epi32(hd.val[ib]); + aux = _mm_cvtepu8_epi16(_mm_unpacklo_epi8(aux, aux)); + auto scales_d = MM256_SET_M128I(aux, aux); +#endif + auto lbits1 = _mm256_loadu_si256((const __m256i *)iq5[ibl].qs+2*ib+0); + auto lbits2 = _mm256_loadu_si256((const __m256i *)iq5[ibl].qs+2*ib+1); + auto hbits128 = _mm_loadu_si128((const __m128i*)iq5[ibl].qh + ib); + auto hbits = MM256_SET_M128I(hbits128, _mm_slli_epi16(hbits128, 4)); + qx[0] = _mm256_or_si256(_mm256_and_si256(lbits1, mf), _mm256_and_si256(m10, hbits)); + qx[1] = _mm256_or_si256(_mm256_and_si256(lbits2, mf), _mm256_and_si256(m10, _mm256_srli_epi16(hbits, 2))); + qx[2] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits1, 4), mf), _mm256_and_si256(m10, _mm256_srli_epi16(hbits, 1))); + qx[3] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits2, 4), mf), _mm256_and_si256(m10, _mm256_srli_epi16(hbits, 3))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = _mm256_loadu_si256((const __m256i*)q8.y[iy][ibl].qs+ib); +#ifdef HAVE_VNNI256 + auto sumi = _mm256_setzero_si256(); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_mullo_epi32(scales_d, sumi)); +#else + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[0], _mm256_shuffle_epi32(y, 0x00)), + _mm256_maddubs_epi16(qx[1], _mm256_shuffle_epi32(y, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[2], _mm256_shuffle_epi32(y, 0xaa)), + _mm256_maddubs_epi16(qx[3], _mm256_shuffle_epi32(y, 0xff))); + // To avoid overflow, we can only add up to 4 q5 x q8 products. + auto sumi = _mm256_add_epi32(_mm256_madd_epi16(scales_d, sumi1), _mm256_madd_epi16(scales_d, sumi2)); + isum[iy] = _mm256_add_epi32(isum[iy], sumi); +#endif + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(isum[iy]), acc[iy]); + isum[iy] = _mm256_setzero_si256(); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm_add_ps(_mm256_castps256_ps128(acc[iy]), _mm256_extractf128_ps(acc[iy], 1)); + acc[iy] = _mm256_setzero_ps(); + info.store(ix+0, iy, sum); + } + } +} + +template +static void mul_mat_q6_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto m4 = _mm256_set1_epi8(0xf); + auto m3 = _mm256_set1_epi8(0x30); + static const uint8_t k_shuff[32] = {0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15, 0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15}; + auto shuff = _mm256_loadu_si256((const __m256i *)k_shuff); +#ifdef HAVE_FANCY_SIMD + __m256i isum[nrc_y]; +#elif !defined(HAVE_VNNI256) + auto m1 = _mm256_set1_epi16(1); +#endif + int nbl = n / QK_K; + __m256 acc[nrc_y] = {}; + __m256i qx[4]; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q6_k_r4 * iq6 = (const block_q6_k_r4 *)((const char *)vx + (ix+0)*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto dl = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq6[ibl].d)); + auto d4 = _mm256_set_m128(dl, dl); +#ifndef HAVE_FANCY_SIMD + if constexpr (nrc_y == 1) { + d4 = _mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(0, ibl))); + } +#endif + { +#ifndef HAVE_FANCY_SIMD + auto min = _mm256_mul_ps(d4, _mm256_set1_ps(-32.f)); +#endif + auto t1 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i *)iq6[ibl].scales+0)), shuff); // blocks 0, 1, 2, 3 for each row + auto t2 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i *)iq6[ibl].scales+1)), shuff); // blocks 4, 5, 6, 7 for each row + auto t3 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i *)iq6[ibl].scales+2)), shuff); // blocks 8, 9, 10, 11 for each row + auto t4 = _mm256_shuffle_epi8(_mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i *)iq6[ibl].scales+3)), shuff); // blocks 12, 13, 14, 15 for each row + auto s1 = MM256_SET_M128I(_mm256_extracti128_si256(t3, 0), _mm256_extracti128_si256(t1, 0)); // blocks 0, 1, 8, 9 + auto s2 = MM256_SET_M128I(_mm256_extracti128_si256(t3, 1), _mm256_extracti128_si256(t1, 1)); // blocks 2, 3, 10, 11 + auto s3 = MM256_SET_M128I(_mm256_extracti128_si256(t4, 0), _mm256_extracti128_si256(t2, 0)); // blocks 4, 5, 12, 13 + auto s4 = MM256_SET_M128I(_mm256_extracti128_si256(t4, 1), _mm256_extracti128_si256(t2, 1)); // blocks 6, 7, 14, 15 +#ifdef HAVE_FANCY_SIMD + s1 = _mm256_mullo_epi16(s1, _mm256_set1_epi16(-32)); + s2 = _mm256_mullo_epi16(s2, _mm256_set1_epi16(-32)); + s3 = _mm256_mullo_epi16(s3, _mm256_set1_epi16(-32)); + s4 = _mm256_mullo_epi16(s4, _mm256_set1_epi16(-32)); +#endif + for (int iy = 0; iy < nrc_y; ++iy) { + auto bsums = q8.load_bsums(iy, ibl); + auto sumi = _mm256_setzero_si256(); +#ifdef HAVE_FANCY_SIMD + sumi = _mm256_dpwssd_epi32(sumi, s1, _mm256_shuffle_epi32(bsums, 0x00)); + sumi = _mm256_dpwssd_epi32(sumi, s2, _mm256_shuffle_epi32(bsums, 0x55)); + sumi = _mm256_dpwssd_epi32(sumi, s3, _mm256_shuffle_epi32(bsums, 0xaa)); + sumi = _mm256_dpwssd_epi32(sumi, s4, _mm256_shuffle_epi32(bsums, 0xff)); + isum[iy] = sumi; +#elif defined(HAVE_VNNI256) + sumi = ggml_mm256_dpwssd_epi32(sumi, s1, _mm256_shuffle_epi32(bsums, 0x00)); + sumi = ggml_mm256_dpwssd_epi32(sumi, s2, _mm256_shuffle_epi32(bsums, 0x55)); + sumi = ggml_mm256_dpwssd_epi32(sumi, s3, _mm256_shuffle_epi32(bsums, 0xaa)); + sumi = ggml_mm256_dpwssd_epi32(sumi, s4, _mm256_shuffle_epi32(bsums, 0xff)); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(min, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(min, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#else + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s1, _mm256_shuffle_epi32(bsums, 0x00))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s2, _mm256_shuffle_epi32(bsums, 0x55))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s3, _mm256_shuffle_epi32(bsums, 0xaa))); + sumi = _mm256_add_epi32(sumi, _mm256_madd_epi16(s4, _mm256_shuffle_epi32(bsums, 0xff))); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(min, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(min, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#endif + } + } + const uint32_t * scales = (const uint32_t *)iq6[ibl].scales; + for (int ib = 0; ib < QK_K/32; ++ib) { + auto iscales = _mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i *)(scales + 2*ib))); +#ifndef HAVE_FANCY_SIMD + auto scales = _mm256_mul_ps(d4, _mm256_cvtepi32_ps(iscales)); +#endif + auto lbits1 = _mm256_loadu_si256((const __m256i *)iq6[ibl].ql+2*ib+0); + auto lbits2 = _mm256_loadu_si256((const __m256i *)iq6[ibl].ql+2*ib+1); + auto hbits = _mm256_loadu_si256((const __m256i *)iq6[ibl].qh+ib); + qx[0] = _mm256_or_si256(_mm256_and_si256(lbits1, m4), _mm256_and_si256(m3, _mm256_slli_epi16(hbits, 4))); + qx[1] = _mm256_or_si256(_mm256_and_si256(lbits2, m4), _mm256_and_si256(m3, hbits)); + qx[2] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits1, 4), m4), _mm256_and_si256(m3, _mm256_slli_epi16(hbits, 2))); + qx[3] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits2, 4), m4), _mm256_and_si256(m3, _mm256_srli_epi16(hbits, 2))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = _mm256_loadu_si256((const __m256i*)q8.y[iy][ibl].qs+ib); +#ifdef HAVE_FANCY_SIMD + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_mullo_epi32(iscales, sumi)); +#elif defined(HAVE_VNNI256) + auto sumi = _mm256_setzero_si256(); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(scales, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(scales, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#else + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[0], _mm256_shuffle_epi32(y, 0x00)), + _mm256_maddubs_epi16(qx[1], _mm256_shuffle_epi32(y, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[2], _mm256_shuffle_epi32(y, 0xaa)), + _mm256_maddubs_epi16(qx[3], _mm256_shuffle_epi32(y, 0xff))); + // Quants are in 0...63, so we can add at most 4 as int16_t to be sure of no int16_t overflow + auto sumi = _mm256_add_epi32(_mm256_madd_epi16(m1, sumi1), _mm256_madd_epi16(m1, sumi2)); + if constexpr (nrc_y == 1) { + acc[iy] = _mm256_fmadd_ps(scales, _mm256_cvtepi32_ps(sumi), acc[iy]); + } else { + acc[iy] = _mm256_fmadd_ps(_mm256_mul_ps(scales, _mm256_set1_ps(q8.scale(iy, ibl))), _mm256_cvtepi32_ps(sumi), acc[iy]); + } +#endif + } + } +#ifdef HAVE_FANCY_SIMD + for (int iy = 0; iy < nrc_y; ++iy) { + auto d4y = _mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(iy, ibl))); + acc[iy] = _mm256_fmadd_ps(d4y, _mm256_cvtepi32_ps(isum[iy]), acc[iy]); + } +#endif + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm_add_ps(_mm256_castps256_ps128(acc[iy]), _mm256_extractf128_ps(acc[iy], 1)); + acc[iy] = _mm256_setzero_ps(); + info.store(ix+0, iy, sum); + } + } +} + +template void set_functions(std::array& funcs) { +#ifdef HAVE_FANCY_SIMD + if constexpr (std::is_same_v) { + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_iqX_k_q8_K_AVX512, Dequantizer, funcs) + } else { + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_AVX512, Dequantizer, funcs) + funcs[0] = mul_mat_qX_K_q8_K_AVX512_1; + } +#else + if constexpr (std::is_same_v || + std::is_same_v || + std::is_same_v) { + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qY_K_q8_K_T, Dequantizer, funcs) + } else { + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, Dequantizer, funcs) + } +#endif +} + +// The HAVE_FANCY_SIMD should only be #if defined(__AVX512_VNNI__ && defined(__AVX512VL__) +template +static void mul_mat_q8_k_r8_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); +#ifndef HAVE_VNNI256 + auto m1 = _mm256_set1_epi16(1); +#endif + int nbl = n / QK_K; + __m256 acc[nrc_y] = {}; + __m256i isum[nrc_y] = {}; + __m256i qx[4]; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q8_k_r8 * iq8 = (const block_q8_k_r8 *)((const char *)vx + (ix+0)*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto d4 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[ibl].d)); + for (int ib = 0; ib < QK_K/16; ++ib) { + qx[0] = _mm256_loadu_si256((const __m256i *)iq8[ibl].qs+4*ib+0); + qx[1] = _mm256_loadu_si256((const __m256i *)iq8[ibl].qs+4*ib+1); + qx[2] = _mm256_loadu_si256((const __m256i *)iq8[ibl].qs+4*ib+2); + qx[3] = _mm256_loadu_si256((const __m256i *)iq8[ibl].qs+4*ib+3); + auto s0 = _mm256_sign_epi8(qx[0], qx[0]); + auto s1 = _mm256_sign_epi8(qx[1], qx[1]); + auto s2 = _mm256_sign_epi8(qx[2], qx[2]); + auto s3 = _mm256_sign_epi8(qx[3], qx[3]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y128 = _mm_loadu_si128((const __m128i*)q8.y[iy][ibl].qs+ib); + auto y = MM256_SET_M128I(y128, y128); +#ifdef HAVE_VNNI256 + isum[iy] = ggml_mm256_dpbusd_epi32(isum[iy], s0, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0])); + isum[iy] = ggml_mm256_dpbusd_epi32(isum[iy], s1, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1])); + isum[iy] = ggml_mm256_dpbusd_epi32(isum[iy], s2, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2])); + isum[iy] = ggml_mm256_dpbusd_epi32(isum[iy], s3, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qx[3])); +#else + auto sumi1 = _mm256_madd_epi16(m1, _mm256_maddubs_epi16(s0, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0]))); + auto sumi2 = _mm256_madd_epi16(m1, _mm256_maddubs_epi16(s1, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1]))); + auto sumi3 = _mm256_madd_epi16(m1, _mm256_maddubs_epi16(s2, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2]))); + auto sumi4 = _mm256_madd_epi16(m1, _mm256_maddubs_epi16(s3, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qx[3]))); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_add_epi32(sumi1, sumi2)); + isum[iy] = _mm256_add_epi32(isum[iy], _mm256_add_epi32(sumi3, sumi4)); +#endif + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto d4y = _mm256_mul_ps(d4, _mm256_set1_ps(q8.scale(iy, ibl))); + acc[iy] = _mm256_fmadd_ps(d4y, _mm256_cvtepi32_ps(isum[iy]), acc[iy]); + isum[iy] = _mm256_setzero_si256(); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = _mm256_setzero_ps(); + } + } +} + +#ifdef HAVE_FANCY_SIMD +template +static void mul_mat_q8_k_r16_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%16 == 0); + Q8 q8(info); + int nbl = n / QK_K; + __m512 acc[nrc_y] = {}; + __m512i isum[nrc_y] = {}; + __m512i qx[4]; + for (int ix = 0; ix < nrc_x; ix += 16) { + const block_q8_k_r16 * iq16 = (const block_q8_k_r16 *)((const char *)vx + ix*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { // Block of 256 + auto d4 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)iq16[ibl].d)); + for (int ib = 0; ib < QK_K/16; ++ib) { + qx[0] = _mm512_loadu_si512((const __m512i *)iq16[ibl].qs+4*ib+0); + qx[1] = _mm512_loadu_si512((const __m512i *)iq16[ibl].qs+4*ib+1); + qx[2] = _mm512_loadu_si512((const __m512i *)iq16[ibl].qs+4*ib+2); + qx[3] = _mm512_loadu_si512((const __m512i *)iq16[ibl].qs+4*ib+3); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y128 = _mm_loadu_si128((const __m128i*)q8.y[iy][ibl].qs+ib); + auto y256 = MM256_SET_M128I(y128, y128); + auto y = _mm512_inserti32x8(_mm512_castsi256_si512(y256), y256, 1); + isum[iy] = _mm512_dpbusd_epi32(isum[iy], qx[0], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + isum[iy] = _mm512_dpbusd_epi32(isum[iy], qx[1], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + isum[iy] = _mm512_dpbusd_epi32(isum[iy], qx[2], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + isum[iy] = _mm512_dpbusd_epi32(isum[iy], qx[3], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + } + } + auto m4 = _mm512_mul_ps(d4, _mm512_set1_ps(-128.f)); + for (int iy = 0; iy < nrc_y; ++iy) { + auto d4y = _mm512_mul_ps(d4, _mm512_set1_ps(q8.scale(iy, ibl))); + acc[iy] = _mm512_fmadd_ps(d4y, _mm512_cvtepi32_ps(isum[iy]), acc[iy]); + acc[iy] = _mm512_fmadd_ps(m4, _mm512_set1_ps(q8.y[iy][ibl].sum), acc[iy]); + isum[iy] = _mm512_setzero_si512(); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = _mm512_setzero_ps(); + } + } +} +#endif + +template +static void mul_mat_q8_KV_q8_KV(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + GGML_ASSERT(n%32 == 0); + __m256i qx[4]; +#ifndef HAVE_FANCY_SIMD + __m256i sx[4]; + auto m1 = _mm256_set1_epi16(1); +#endif + __m256i acc[nrc_y] = {}; + float dy[nrc_y]; +#ifdef HAVE_FANCY_SIMD + int32_t sy[nrc_y]; +#endif + const int8_t * q8y[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) { + auto dptr = (const float *)info.src1_row(iy); + dy[iy] = dptr[0]; +#ifdef HAVE_FANCY_SIMD + auto iptr = (const int32_t *)(dptr + 1); + sy[iy] = -127*iptr[0]; +#endif + q8y[iy] = (const int8_t *)(dptr + 2); + } + const int8_t * q8x[4]; + float dx[4]; + for (int ix = 0; ix < nrc_x; ix += 4) { + for (int kx = 0; kx < 4; ++kx) { + auto dptr = (const float *)((const char *)vx + (ix+kx)*bx); + dx[kx] = dptr[0]; + q8x[kx] = (const int8_t *)(dptr + 2); + } + for (int i = 0; i < n/32; ++i) { + for (int kx = 0; kx < 4; ++kx) qx[kx] = _mm256_loadu_si256((const __m256i *)q8x[kx] + i); + auto t0 = _mm256_unpacklo_epi32(qx[0], qx[1]); + auto t1 = _mm256_unpacklo_epi32(qx[2], qx[3]); + auto t2 = _mm256_unpackhi_epi32(qx[0], qx[1]); + auto t3 = _mm256_unpackhi_epi32(qx[2], qx[3]); +#ifdef HAVE_FANCY_SIMD + qx[0] = _mm256_add_epi8(_mm256_unpacklo_epi64(t0, t1), _mm256_set1_epi8(127)); + qx[1] = _mm256_add_epi8(_mm256_unpackhi_epi64(t0, t1), _mm256_set1_epi8(127)); + qx[2] = _mm256_add_epi8(_mm256_unpacklo_epi64(t2, t3), _mm256_set1_epi8(127)); + qx[3] = _mm256_add_epi8(_mm256_unpackhi_epi64(t2, t3), _mm256_set1_epi8(127)); +#else + qx[0] = _mm256_unpacklo_epi64(t0, t1); sx[0] = _mm256_sign_epi8(qx[0], qx[0]); + qx[1] = _mm256_unpackhi_epi64(t0, t1); sx[1] = _mm256_sign_epi8(qx[1], qx[1]); + qx[2] = _mm256_unpacklo_epi64(t2, t3); sx[2] = _mm256_sign_epi8(qx[2], qx[2]); + qx[3] = _mm256_unpackhi_epi64(t2, t3); sx[3] = _mm256_sign_epi8(qx[3], qx[3]); +#endif + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = _mm256_loadu_si256((const __m256i *)q8y[iy] + i); +#ifdef HAVE_FANCY_SIMD + acc[iy] = _mm256_dpbusd_epi32(acc[iy], qx[0], _mm256_shuffle_epi32(y, 0x00)); + acc[iy] = _mm256_dpbusd_epi32(acc[iy], qx[1], _mm256_shuffle_epi32(y, 0x55)); + acc[iy] = _mm256_dpbusd_epi32(acc[iy], qx[2], _mm256_shuffle_epi32(y, 0xaa)); + acc[iy] = _mm256_dpbusd_epi32(acc[iy], qx[3], _mm256_shuffle_epi32(y, 0xff)); +#else + auto dot1 = _mm256_maddubs_epi16(sx[0], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0])); + auto dot2 = _mm256_maddubs_epi16(sx[1], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1])); + auto dot3 = _mm256_maddubs_epi16(sx[2], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2])); + auto dot4 = _mm256_maddubs_epi16(sx[3], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qx[3])); + auto dot12 = _mm256_add_epi32(_mm256_madd_epi16(m1, dot1), _mm256_madd_epi16(m1, dot2)); + auto dot34 = _mm256_add_epi32(_mm256_madd_epi16(m1, dot3), _mm256_madd_epi16(m1, dot4)); + acc[iy] = _mm256_add_epi32(acc[iy], _mm256_add_epi32(dot12, dot34)); +#endif + } + } + auto scales_x = _mm_loadu_ps(dx); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = _mm_add_epi32(_mm256_castsi256_si128(acc[iy]), _mm256_extracti128_si256(acc[iy], 1)); +#ifdef HAVE_FANCY_SIMD + sumi = _mm_add_epi32(sumi, _mm_set1_epi32(sy[iy])); +#endif + auto scale = _mm_mul_ps(scales_x, _mm_set1_ps(dy[iy])); + info.store(ix, iy, _mm_mul_ps(scale, _mm_cvtepi32_ps(sumi))); + acc[iy] = _mm256_setzero_si256(); + } + } +} + +// The HAVE_FANCY_SIMD should only be #if defined(__AVX512_VNNI__ && defined(__AVX512VL__) +template +static void mul_mat_q8_KV_r8_q8_KV(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(n%32 == 0); + GGML_ASSERT(nrc_x%8 == 0); +#ifndef HAVE_FANCY_SIMD + auto m1 = _mm256_set1_epi16(1); +#endif + int nb = n / 16; + __m256i acc[nrc_y] = {}; + __m256i qx[4]; + float dy[nrc_y]; +#ifdef HAVE_FANCY_SIMD + float sy[nrc_y]; +#endif + const int8_t * q8y[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) { + auto dptr = (const float *)info.src1_row(iy); + dy[iy] = dptr[0]; +#ifdef HAVE_FANCY_SIMD + auto iptr = (const int32_t *)(dptr + 1); + sy[iy] = -127*iptr[0]; +#endif + q8y[iy] = (const int8_t *)(dptr + 2); + } + for (int ix = 0; ix < nrc_x; ix += 8) { + auto dptr = (const float *)((const char *)vx + ix*bx); + auto dx = _mm256_loadu_ps(dptr); + auto q8x = (const int8_t *)(dptr + 8); + for (int ib = 0; ib < nb; ++ib) { // Blocks of 16 for 8 interleaved rows + qx[0] = _mm256_loadu_si256((const __m256i *)q8x+4*ib+0); + qx[1] = _mm256_loadu_si256((const __m256i *)q8x+4*ib+1); + qx[2] = _mm256_loadu_si256((const __m256i *)q8x+4*ib+2); + qx[3] = _mm256_loadu_si256((const __m256i *)q8x+4*ib+3); +#ifndef HAVE_FANCY_SIMD + auto s0 = _mm256_sign_epi8(qx[0], qx[0]); + auto s1 = _mm256_sign_epi8(qx[1], qx[1]); + auto s2 = _mm256_sign_epi8(qx[2], qx[2]); + auto s3 = _mm256_sign_epi8(qx[3], qx[3]); +#else + qx[0] = _mm256_add_epi8(qx[0], _mm256_set1_epi8(127)); + qx[1] = _mm256_add_epi8(qx[1], _mm256_set1_epi8(127)); + qx[2] = _mm256_add_epi8(qx[2], _mm256_set1_epi8(127)); + qx[3] = _mm256_add_epi8(qx[3], _mm256_set1_epi8(127)); +#endif + for (int iy = 0; iy < nrc_y; ++iy) { + auto y128 = _mm_loadu_si128((const __m128i*)q8y[iy]+ib); + auto y = MM256_SET_M128I(y128, y128); +#ifdef HAVE_FANCY_SIMD + acc[iy] = _mm256_dpbusd_epi32(acc[iy], qx[0], _mm256_shuffle_epi32(y, 0x00)); + acc[iy] = _mm256_dpbusd_epi32(acc[iy], qx[1], _mm256_shuffle_epi32(y, 0x55)); + acc[iy] = _mm256_dpbusd_epi32(acc[iy], qx[2], _mm256_shuffle_epi32(y, 0xaa)); + acc[iy] = _mm256_dpbusd_epi32(acc[iy], qx[3], _mm256_shuffle_epi32(y, 0xff)); +#else + auto sumi1 = _mm256_maddubs_epi16(s0, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0])); + auto sumi2 = _mm256_maddubs_epi16(s1, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1])); + auto sumi3 = _mm256_maddubs_epi16(s2, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2])); + auto sumi4 = _mm256_maddubs_epi16(s3, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qx[3])); + auto sumi12 = _mm256_add_epi32(_mm256_madd_epi16(m1, sumi1), _mm256_madd_epi16(m1, sumi2)); + auto sumi34 = _mm256_add_epi32(_mm256_madd_epi16(m1, sumi3), _mm256_madd_epi16(m1, sumi4)); + acc[iy] = _mm256_add_epi32(acc[iy], _mm256_add_epi32(sumi12, sumi34)); +#endif + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto scale = _mm256_mul_ps(dx, _mm256_set1_ps(dy[iy])); +#ifdef HAVE_FANCY_SIMD + acc[iy] = _mm256_add_epi32(acc[iy], _mm256_set1_epi32(sy[iy])); +#endif + info.store(ix, iy, _mm256_mul_ps(scale, _mm256_cvtepi32_ps(acc[iy]))); + acc[iy] = _mm256_setzero_si256(); + } + } +} + +typedef struct { + ggml_half d[16]; + int8_t qs[8*QK8_1]; +} block_q8_1_r8; + +void iqk_convert_q2_k_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { +#ifdef HAVE_FANCY_SIMD + constexpr int k_nr = 16; + using block_q8_k_r = block_q8_k_r16; +#else + constexpr int k_nr = 8; + using block_q8_k_r = block_q8_k_r8; +#endif + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%k_nr == 0); + + int nb = n/QK_K; + + const block_q2_K * x8[k_nr]; + + block_q8_k_r * y = (block_q8_k_r *)vy; + + float f_values[QK_K]; + uint32_t block[8]; + + __m256i xv[4]; + + auto ml = _mm256_set1_epi8(0x03); + auto sign_bit = _mm256_set1_ps(-0.0f); + auto perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7); + + for (int ix = 0; ix < nrc_x; ix += k_nr) { + for (int k = 0; k < k_nr; ++k) x8[k] = (const block_q2_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < k_nr; ++k) { + auto vd = _mm256_set1_ps(GGML_FP16_TO_FP32(x8[k][i].d)); + auto vm = _mm256_mul_ps(_mm256_set1_ps(GGML_FP16_TO_FP32(x8[k][i].dmin)), _mm256_set1_ps(-1.f)); + auto block_max = _mm256_setzero_ps(); + for (int i128 = 0; i128 < 2; ++i128) { + auto bits = _mm256_loadu_si256((const __m256i *)x8[k][i].qs+i128); + xv[0] = _mm256_and_si256(bits, ml); + xv[1] = _mm256_and_si256(_mm256_srli_epi16(bits, 2), ml); + xv[2] = _mm256_and_si256(_mm256_srli_epi16(bits, 4), ml); + xv[3] = _mm256_and_si256(_mm256_srli_epi16(bits, 6), ml); + for (int l = 0; l < 4; ++l) { + auto q1 = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(xv[l])); + auto q2 = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(xv[l], 1)); + q1 = _mm256_mullo_epi16(q1, _mm256_set1_epi16(x8[k][i].scales[8*i128 + 2*l + 0] & 0xf)); + q2 = _mm256_mullo_epi16(q2, _mm256_set1_epi16(x8[k][i].scales[8*i128 + 2*l + 1] & 0xf)); + auto m1 = _mm256_mul_ps(vm, _mm256_set1_ps(x8[k][i].scales[8*i128 + 2*l + 0] >> 4)); + auto m2 = _mm256_mul_ps(vm, _mm256_set1_ps(x8[k][i].scales[8*i128 + 2*l + 1] >> 4)); + auto v0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_castsi256_si128(q1))), vd, m1); + auto v1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(q1, 1))), vd, m1); + auto v2 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_castsi256_si128(q2))), vd, m2); + auto v3 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(q2, 1))), vd, m2); + auto max = _mm256_max_ps(_mm256_max_ps(_mm256_andnot_ps(sign_bit, v0), _mm256_andnot_ps(sign_bit, v1)), + _mm256_max_ps(_mm256_andnot_ps(sign_bit, v2), _mm256_andnot_ps(sign_bit, v3))); + block_max = _mm256_max_ps(block_max, max); + _mm256_storeu_ps(f_values + 128*i128 + 32*l + 0, v0); + _mm256_storeu_ps(f_values + 128*i128 + 32*l + 8, v1); + _mm256_storeu_ps(f_values + 128*i128 + 32*l + 16, v2); + _mm256_storeu_ps(f_values + 128*i128 + 32*l + 24, v3); + } + } + auto max4 = _mm_max_ps(_mm256_extractf128_ps(block_max, 1), _mm256_castps256_ps128(block_max)); + max4 = _mm_max_ps(max4, _mm_movehl_ps(max4, max4)); + max4 = _mm_max_ss(max4, _mm_movehdup_ps(max4)); + float d = _mm_cvtss_f32(max4)/127.f; + auto id = _mm256_set1_ps(d != 0.0f ? 1/d : 0.0f); + y[i].d[k] = GGML_FP32_TO_FP16(d); + for (int ib32 = 0; ib32 < 8; ++ib32) { + auto v0 = _mm256_loadu_ps(f_values + 32*ib32 + 0); + auto v1 = _mm256_loadu_ps(f_values + 32*ib32 + 8); + auto v2 = _mm256_loadu_ps(f_values + 32*ib32 + 16); + auto v3 = _mm256_loadu_ps(f_values + 32*ib32 + 24); + auto i0 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(v0, id), _MM_ROUND_NEAREST)); + auto i1 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(v1, id), _MM_ROUND_NEAREST)); + auto i2 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(v2, id), _MM_ROUND_NEAREST)); + auto i3 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(v3, id), _MM_ROUND_NEAREST)); + i0 = _mm256_packs_epi32(i0, i1); + i2 = _mm256_packs_epi32(i2, i3); + i0 = _mm256_packs_epi16(i0, i2); + i0 = _mm256_permutevar8x32_epi32(i0, perm); + + _mm256_storeu_si256((__m256i *)block, i0); + auto q8 = (uint32_t *)y[i].qs + 8*k_nr*ib32; + for (int l = 0; l < 8; ++l) { + q8[k_nr*l + k] = block[l]; + } + } + } +#ifdef HAVE_FANCY_SIMD + for (int l = 0; l < 64; ++l) { + auto v = _mm512_xor_si512(_mm512_loadu_si512((const __m512i *)y[i].qs + l), _mm512_set1_epi8(-128)); + _mm512_storeu_si512((__m512i *)y[i].qs + l, v); + } +#endif + } + y += nb; + } +} + +void iqk_convert_q4_k_q8_1_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q4_K * x8[8]; + + block_q8_1_r8 * y = (block_q8_1_r8 *)vy; + + ggml_half dh[16]; + uint16_t all_ls[128]; + + uint32_t utmp[4]; + const uint8_t * u8 = (const uint8_t *)utmp; + uint32_t block[8]; + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q4_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + dh[k+0] = x8[k][i].d; + dh[k+8] = x8[k][i].dmin; + make_q4_scales(x8[k][i].scales, utmp); + auto qs = x8[k][i].qs; + for (int ib64 = 0; ib64 < 4; ++ib64) { + all_ls[8*(2*ib64 + 0) + k ] = u8[2*ib64+0]; + all_ls[8*(2*ib64 + 1) + k ] = u8[2*ib64+1]; + all_ls[8*(2*ib64 + 0) + k + 64] = u8[2*ib64+8]; + all_ls[8*(2*ib64 + 1) + k + 64] = u8[2*ib64+9]; + auto bits = _mm256_loadu_si256((const __m256i *)qs+ib64); + auto values1 = _mm256_and_si256(bits, _mm256_set1_epi8(0xf)); + auto values2 = _mm256_and_si256(_mm256_srli_epi16(bits, 4), _mm256_set1_epi8(0xf)); + _mm256_storeu_si256((__m256i *)block, values1); + auto q8 = (uint32_t *)y[2*ib64+0].qs; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + _mm256_storeu_si256((__m256i *)block, values2); + q8 = (uint32_t *)y[2*ib64+1].qs; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + } + } + auto vd = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)dh+0)); + auto vm = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)dh+1)); + vm = _mm256_mul_ps(_mm256_set1_ps(-1.f), vm); + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + auto iscales16 = _mm_loadu_si128((const __m128i *)all_ls + ib32); + auto iscales32 = _mm256_cvtepi16_epi32(iscales16); + auto scales = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(iscales32)); + _mm_storeu_si128((__m128i *)y[ib32].d+0, _mm256_cvtps_ph(scales, _MM_FROUND_TO_NEAREST_INT)); + iscales16 = _mm_loadu_si128((const __m128i *)all_ls + ib32 + 8); + iscales32 = _mm256_cvtepi16_epi32(iscales16); + scales = _mm256_mul_ps(vm, _mm256_cvtepi32_ps(iscales32)); + _mm_storeu_si128((__m128i *)y[ib32].d+1, _mm256_cvtps_ph(scales, _MM_FROUND_TO_NEAREST_INT)); + } + y += QK_K/32; + } + } +} + +void iqk_convert_q5_k_q8_1_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q5_K * x8[8]; + + block_q8_1_r8 * y = (block_q8_1_r8 *)vy; + + ggml_half dh[16]; + uint16_t all_ls[128]; + + uint32_t utmp[4]; + const uint8_t * u8 = (const uint8_t *)utmp; + uint32_t block[8]; + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q5_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + dh[k+0] = x8[k][i].d; + dh[k+8] = x8[k][i].dmin; + make_q4_scales(x8[k][i].scales, utmp); + auto qs = x8[k][i].qs; + auto hbits = _mm256_loadu_si256((const __m256i *)x8[k][i].qh); + for (int ib64 = 0; ib64 < 4; ++ib64) { + all_ls[8*(2*ib64 + 0) + k ] = u8[2*ib64+0]; + all_ls[8*(2*ib64 + 1) + k ] = u8[2*ib64+1]; + all_ls[8*(2*ib64 + 0) + k + 64] = u8[2*ib64+8]; + all_ls[8*(2*ib64 + 1) + k + 64] = u8[2*ib64+9]; + auto bits = _mm256_loadu_si256((const __m256i *)qs+ib64); + auto values1 = _mm256_and_si256(bits, _mm256_set1_epi8(0xf)); + auto values2 = _mm256_and_si256(_mm256_srli_epi16(bits, 4), _mm256_set1_epi8(0xf)); + values1 = _mm256_or_si256(values1, _mm256_and_si256(_mm256_set1_epi8(0x10), _mm256_slli_epi16(hbits, 4))); + values2 = _mm256_or_si256(values2, _mm256_and_si256(_mm256_set1_epi8(0x10), _mm256_slli_epi16(hbits, 3))); + hbits = _mm256_srli_epi16(hbits, 2); + _mm256_storeu_si256((__m256i *)block, values1); + auto q8 = (uint32_t *)y[2*ib64+0].qs; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + _mm256_storeu_si256((__m256i *)block, values2); + q8 = (uint32_t *)y[2*ib64+1].qs; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + } + } + auto vd = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)dh+0)); + auto vm = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)dh+1)); + vm = _mm256_mul_ps(_mm256_set1_ps(-1.f), vm); + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + auto iscales16 = _mm_loadu_si128((const __m128i *)all_ls + ib32); + auto iscales32 = _mm256_cvtepi16_epi32(iscales16); + auto scales = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(iscales32)); + _mm_storeu_si128((__m128i *)y[ib32].d+0, _mm256_cvtps_ph(scales, _MM_FROUND_TO_NEAREST_INT)); + iscales16 = _mm_loadu_si128((const __m128i *)all_ls + ib32 + 8); + iscales32 = _mm256_cvtepi16_epi32(iscales16); + scales = _mm256_mul_ps(vm, _mm256_cvtepi32_ps(iscales32)); + _mm_storeu_si128((__m128i *)y[ib32].d+1, _mm256_cvtps_ph(scales, _MM_FROUND_TO_NEAREST_INT)); + } + y += QK_K/32; + } + } +} + +void iqk_convert_q6_k_q8_0_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q6_K * x8[8]; + + block_q8_0_r8 * y = (block_q8_0_r8 *)vy; + + float all_s[64]; + uint32_t block[8]; + __m256i values[8]; + + auto ml = _mm256_set1_epi8(0x0f); + auto mh = _mm256_set1_epi8(0x30); + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q6_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + float d = GGML_FP16_TO_FP32(x8[k][i].d); + auto ql = x8[k][i].ql; + auto qh = x8[k][i].qh; + for (int i128 = 0; i128 < 2; ++i128) { + auto lbits1 = _mm256_loadu_si256((const __m256i *)ql + 2*i128 + 0); + auto lbits2 = _mm256_loadu_si256((const __m256i *)ql + 2*i128 + 1); + auto hbits = _mm256_loadu_si256((const __m256i *)qh + i128); + values[4*i128+0] = _mm256_or_si256(_mm256_and_si256(lbits1, ml), _mm256_and_si256(_mm256_slli_epi16(hbits, 4), mh)); + values[4*i128+1] = _mm256_or_si256(_mm256_and_si256(lbits2, ml), _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + values[4*i128+2] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits1, 4), ml), _mm256_and_si256(hbits, mh)); + values[4*i128+3] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(lbits2, 4), ml), _mm256_and_si256(_mm256_srli_epi16(hbits, 2), mh)); + } + for (int ib32 = 0; ib32 < 8; ++ib32) { + // We have two blocks of 16 with different scales + // We multiply the quants with the scales, find the max value, and convert to 8-bit quants with a single block scale. + auto q8 = _mm256_add_epi8(values[ib32], _mm256_set1_epi8(-32)); + auto q16_l = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8)); + auto q16_h = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8, 1)); + q16_l = _mm256_mullo_epi16(q16_l, _mm256_set1_epi16(x8[k][i].scales[2*ib32+0])); + q16_h = _mm256_mullo_epi16(q16_h, _mm256_set1_epi16(x8[k][i].scales[2*ib32+1])); + auto abs_q16_l = _mm256_sign_epi16(q16_l, q16_l); + auto abs_q16_h = _mm256_sign_epi16(q16_h, q16_h); + auto max_q16 = _mm256_max_epi16(abs_q16_l, abs_q16_h); + auto max_q32 = _mm256_cvtepi16_epi32(_mm_max_epi16(_mm256_castsi256_si128(max_q16), _mm256_extracti128_si256(max_q16, 1))); + auto imax4 = _mm_max_epi32(_mm256_castsi256_si128(max_q32), _mm256_extracti128_si256(max_q32, 1)); + auto max4 = _mm_cvtepi32_ps(imax4); + max4 = _mm_max_ps( max4, _mm_movehl_ps( max4, max4 ) ); + max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4 ) ); + float max = _mm_cvtss_f32(max4) / 127; + all_s[8*ib32+k] = d*max; + if (max > 1e-9f) { + auto scale = _mm256_set1_ps(1/max); + auto i0 = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16_l)); + auto i1 = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16_l, 1)); + auto i2 = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16_h)); + auto i3 = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16_h, 1)); + i0 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i0)), _MM_ROUND_NEAREST)); + i1 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i1)), _MM_ROUND_NEAREST)); + i2 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i2)), _MM_ROUND_NEAREST)); + i3 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i3)), _MM_ROUND_NEAREST)); + i0 = _mm256_packs_epi32(i0, i1); + i2 = _mm256_packs_epi32(i2, i3); + i0 = _mm256_packs_epi16(i0, i2); + i0 = _mm256_permutevar8x32_epi32(i0, _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7)); + _mm256_storeu_si256((__m256i *)block, i0); + } else { + _mm256_storeu_si256((__m256i *)block, _mm256_setzero_si256()); + } + auto qs = (uint32_t *)y[ib32].qs; + for (int l = 0; l < 4; ++l) { + qs[8*l + k + 0] = block[l + 0]; + qs[8*l + k + 32] = block[l + 4]; + } + } + } + for (int ib32 = 0; ib32 < 8; ++ib32) { + _mm_storeu_si128((__m128i *)y[ib32].d, _mm256_cvtps_ph(_mm256_loadu_ps(all_s + 8*ib32), _MM_FROUND_TO_NEAREST_INT)); + } + y += QK_K/32; + } + } +} + +void iqk_convert_q3_k_q8_0_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q3_K * x8[8]; + + block_q8_0_r8 * y = (block_q8_0_r8 *)vy; + + float all_s[64]; + uint32_t block[8]; + __m256i values[8]; + + ScaleQ3 sc3; + auto ml = _mm256_set1_epi8(0x03); + auto mh = _mm256_set1_epi8(0x04); + + union { __m256i vec; int16_t val[16]; } helper; + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q3_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + float d = GGML_FP16_TO_FP32(x8[k][i].d); + auto hbits = _mm256_loadu_si256((const __m256i *)x8[k][i].hmask); + for (int i128 = 0; i128 < 2; ++i128) { + auto q2bits = _mm256_loadu_si256((const __m256i *)x8[k][i].qs + i128); + values[4*i128+0] = _mm256_and_si256(q2bits, ml); + values[4*i128+1] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 2), ml); + values[4*i128+2] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 4), ml); + values[4*i128+3] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 6), ml); + values[4*i128+0] = _mm256_or_si256(values[4*i128+0], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + values[4*i128+1] = _mm256_or_si256(values[4*i128+1], _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh)); + values[4*i128+2] = _mm256_or_si256(values[4*i128+2], _mm256_and_si256(hbits, mh)); + values[4*i128+3] = _mm256_or_si256(values[4*i128+3], _mm256_and_si256(_mm256_srli_epi16(hbits, 1), mh)); + values[4*i128+0] = _mm256_sub_epi8(values[4*i128+0], mh); + values[4*i128+1] = _mm256_sub_epi8(values[4*i128+1], mh); + values[4*i128+2] = _mm256_sub_epi8(values[4*i128+2], mh); + values[4*i128+3] = _mm256_sub_epi8(values[4*i128+3], mh); + hbits = _mm256_srli_epi16(hbits, 4); + } + helper.vec = _mm256_cvtepi8_epi16(sc3.make_scales((const uint16_t *)x8[k][i].scales)); + for (int ib32 = 0; ib32 < 8; ++ib32) { + auto q16_l = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(values[ib32])); + auto q16_h = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(values[ib32], 1)); + q16_l = _mm256_mullo_epi16(q16_l, _mm256_set1_epi16(helper.val[2*ib32+0])); + q16_h = _mm256_mullo_epi16(q16_h, _mm256_set1_epi16(helper.val[2*ib32+1])); + auto abs_q16_l = _mm256_sign_epi16(q16_l, q16_l); + auto abs_q16_h = _mm256_sign_epi16(q16_h, q16_h); + auto max_q16 = _mm256_max_epi16(abs_q16_l, abs_q16_h); + auto max_q32 = _mm256_cvtepi16_epi32(_mm_max_epi16(_mm256_castsi256_si128(max_q16), _mm256_extracti128_si256(max_q16, 1))); + auto imax4 = _mm_max_epi32(_mm256_castsi256_si128(max_q32), _mm256_extracti128_si256(max_q32, 1)); + auto max4 = _mm_cvtepi32_ps(imax4); + max4 = _mm_max_ps( max4, _mm_movehl_ps( max4, max4 ) ); + max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4 ) ); + float max = _mm_cvtss_f32(max4) / 127; + all_s[8*ib32+k] = d*max; + if (max > 1e-9f) { + auto scale = _mm256_set1_ps(1/max); + auto i0 = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16_l)); + auto i1 = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16_l, 1)); + auto i2 = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16_h)); + auto i3 = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16_h, 1)); + i0 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i0)), _MM_ROUND_NEAREST)); + i1 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i1)), _MM_ROUND_NEAREST)); + i2 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i2)), _MM_ROUND_NEAREST)); + i3 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i3)), _MM_ROUND_NEAREST)); + i0 = _mm256_packs_epi32(i0, i1); + i2 = _mm256_packs_epi32(i2, i3); + i0 = _mm256_packs_epi16(i0, i2); + i0 = _mm256_permutevar8x32_epi32(i0, _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7)); + _mm256_storeu_si256((__m256i *)block, i0); + } else { + _mm256_storeu_si256((__m256i *)block, _mm256_setzero_si256()); + } + auto qs = (uint32_t *)y[ib32].qs; + for (int l = 0; l < 4; ++l) { + qs[8*l + k + 0] = block[l + 0]; + qs[8*l + k + 32] = block[l + 4]; + } + } + } + for (int ib32 = 0; ib32 < 8; ++ib32) { + _mm_storeu_si128((__m128i *)y[ib32].d, _mm256_cvtps_ph(_mm256_loadu_ps(all_s + 8*ib32), _MM_FROUND_TO_NEAREST_INT)); + } + y += QK_K/32; + } + } +} + +void iqk_convert_q3_k_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { +#ifdef HAVE_FANCY_SIMD + constexpr int k_nr = 16; + using block_q8_k_r = block_q8_k_r16; +#else + constexpr int k_nr = 8; + using block_q8_k_r = block_q8_k_r8; +#endif + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%k_nr == 0); + + int nb = n/QK_K; + + const block_q3_K * x8[k_nr]; + + block_q8_k_r * y = (block_q8_k_r *)vy; + + uint32_t block[8]; + __m256i values[8]; + + ScaleQ3 sc3; + auto ml = _mm256_set1_epi8(0x03); + auto mh = _mm256_set1_epi8(0x04); + + union { __m256i vec; int16_t val[16]; } helper; + + for (int ix = 0; ix < nrc_x; ix += k_nr) { + for (int k = 0; k < k_nr; ++k) x8[k] = (const block_q3_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < k_nr; ++k) { + float d = GGML_FP16_TO_FP32(x8[k][i].d); + auto hbits = _mm256_loadu_si256((const __m256i *)x8[k][i].hmask); + helper.vec = _mm256_cvtepi8_epi16(sc3.make_scales((const uint16_t *)x8[k][i].scales)); + auto max_i16 = _mm256_setzero_si256(); + for (int i128 = 0; i128 < 2; ++i128) { + auto q2bits = _mm256_loadu_si256((const __m256i *)x8[k][i].qs + i128); + values[4*i128+0] = _mm256_and_si256(q2bits, ml); + values[4*i128+1] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 2), ml); + values[4*i128+2] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 4), ml); + values[4*i128+3] = _mm256_and_si256(_mm256_srli_epi16(q2bits, 6), ml); + values[4*i128+0] = _mm256_or_si256(values[4*i128+0], _mm256_and_si256(_mm256_slli_epi16(hbits, 2), mh)); + values[4*i128+1] = _mm256_or_si256(values[4*i128+1], _mm256_and_si256(_mm256_slli_epi16(hbits, 1), mh)); + values[4*i128+2] = _mm256_or_si256(values[4*i128+2], _mm256_and_si256(hbits, mh)); + values[4*i128+3] = _mm256_or_si256(values[4*i128+3], _mm256_and_si256(_mm256_srli_epi16(hbits, 1), mh)); + values[4*i128+0] = _mm256_sub_epi8(values[4*i128+0], mh); + values[4*i128+1] = _mm256_sub_epi8(values[4*i128+1], mh); + values[4*i128+2] = _mm256_sub_epi8(values[4*i128+2], mh); + values[4*i128+3] = _mm256_sub_epi8(values[4*i128+3], mh); + hbits = _mm256_srli_epi16(hbits, 4); + + for (int l = 0; l < 4; ++l) { + auto q16_l = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(values[4*i128+l])); + auto q16_h = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(values[4*i128+l], 1)); + q16_l = _mm256_mullo_epi16(_mm256_set1_epi16(helper.val[8*i128+2*l+0]), q16_l); + q16_h = _mm256_mullo_epi16(_mm256_set1_epi16(helper.val[8*i128+2*l+1]), q16_h); + max_i16 = _mm256_max_epi16(max_i16, _mm256_sign_epi16(q16_l, q16_l)); + max_i16 = _mm256_max_epi16(max_i16, _mm256_sign_epi16(q16_h, q16_h)); + } + } + auto max_q32 = _mm256_cvtepi16_epi32(_mm_max_epi16(_mm256_castsi256_si128(max_i16), _mm256_extracti128_si256(max_i16, 1))); + auto imax4 = _mm_max_epi32(_mm256_castsi256_si128(max_q32), _mm256_extracti128_si256(max_q32, 1)); + auto max4 = _mm_cvtepi32_ps(imax4); + max4 = _mm_max_ps(max4, _mm_movehl_ps(max4, max4)); + max4 = _mm_max_ss(max4, _mm_movehdup_ps(max4)); + bool needs_scaling = true; + float dnew = _mm_cvtss_f32(max4) / 127; + if (dnew < 1.f) { + dnew = 1.f; needs_scaling = false; + } + d *= dnew; + y[i].d[k] = GGML_FP32_TO_FP16(d); + auto scale = _mm256_set1_ps(std::abs(dnew) > 1e-9f ? 1/dnew : 0.f); + for (int ib32 = 0; ib32 < 8; ++ib32) { + auto q16_l = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(values[ib32])); + auto q16_h = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(values[ib32], 1)); + q16_l = _mm256_mullo_epi16(q16_l, _mm256_set1_epi16(helper.val[2*ib32+0])); + q16_h = _mm256_mullo_epi16(q16_h, _mm256_set1_epi16(helper.val[2*ib32+1])); + if (needs_scaling) { + auto i0 = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16_l)); + auto i1 = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16_l, 1)); + auto i2 = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(q16_h)); + auto i3 = _mm256_cvtepi16_epi32(_mm256_extracti128_si256(q16_h, 1)); + i0 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i0)), _MM_ROUND_NEAREST)); + i1 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i1)), _MM_ROUND_NEAREST)); + i2 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i2)), _MM_ROUND_NEAREST)); + i3 = _mm256_cvtps_epi32(_mm256_round_ps(_mm256_mul_ps(scale, _mm256_cvtepi32_ps(i3)), _MM_ROUND_NEAREST)); + i0 = _mm256_packs_epi32(i0, i1); + i2 = _mm256_packs_epi32(i2, i3); + i0 = _mm256_packs_epi16(i0, i2); + i0 = _mm256_permutevar8x32_epi32(i0, _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7)); + _mm256_storeu_si256((__m256i *)block, i0); + } else { + // 0, 1, 2, 3, 4, 5, 6, 7, 8, 16, 17, 18, 19, 20, 21, 22, 23, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31 + auto i0 = _mm256_packs_epi16(q16_l, q16_h); + auto i0_l = _mm256_castsi256_si128(i0); + auto i0_h = _mm256_extracti128_si256(i0, 1); + _mm_storeu_si128((__m128i *)block+0, _mm_unpacklo_epi64(i0_l, i0_h)); + _mm_storeu_si128((__m128i *)block+1, _mm_unpackhi_epi64(i0_l, i0_h)); + } + auto qs = (uint32_t *)y[i].qs + 8*k_nr*ib32; + for (int l = 0; l < 8; ++l) { + qs[k_nr*l + k] = block[l]; + } + } + } +#ifdef HAVE_FANCY_SIMD + for (int l = 0; l < 64; ++l) { + auto v = _mm512_xor_si512(_mm512_loadu_si512((const __m512i *)y[i].qs + l), _mm512_set1_epi8(-128)); + _mm512_storeu_si512((__m512i *)y[i].qs + l, v); + } +#endif + } + y += nb; + } +} + +// TODO: move this to iqk_gemm_iquants +void iqk_convert_iq4_xs_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + +#ifdef HAVE_FANCY_SIMD + constexpr int k_nr = 16; + using block_q8_k_r = block_q8_k_r16; +#else + constexpr int k_nr = 8; + using block_q8_k_r = block_q8_k_r8; +#endif + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%k_nr == 0); + + int nb = n/QK_K; + + const block_iq4_xs * x8[k_nr]; + + block_q8_k_r * y = (block_q8_k_r *)vy; + + auto values128 = _mm_loadu_si128((const __m128i *)iq4k_values); + auto values = MM256_SET_M128I(values128, values128); + + int16_t ls[16]; + float dnew[k_nr]; + uint32_t block[8]; + __m256i xv[8]; + + for (int ix = 0; ix < nrc_x; ix += k_nr) { + for (int k = 0; k < k_nr; ++k) x8[k] = (const block_iq4_xs *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < k_nr; ++k) { + float d = GGML_FP16_TO_FP32(x8[k][i].d); + for (int ib32 = 0; ib32 < 8; ++ib32) { + ls[2*ib32+0] = ls[2*ib32+1] = (((x8[k][i].scales_l[ib32/2] >> 4*(ib32%2)) & 0xf) | (((x8[k][i].scales_h >> 2*ib32) & 3) << 4)) - 32; + auto bits = _mm_loadu_si128((const __m128i *)x8[k][i].qs + ib32); + xv[ib32] = _mm256_and_si256(MM256_SET_M128I(_mm_srli_epi16(bits, 4), bits), _mm256_set1_epi8(0xf)); + xv[ib32] = _mm256_shuffle_epi8(values, xv[ib32]); + } + dnew[k] = d * convert_to_q8_k_r8(k, 1.f/127, xv, ls, block, y[i].qs); + } +#ifdef HAVE_FANCY_SIMD + _mm256_storeu_si256((__m256i *)y[i].d, _mm512_cvtps_ph(_mm512_loadu_ps(dnew), _MM_ROUND_NEAREST)); + for (int l = 0; l < 64; ++l) { + auto v = _mm512_xor_si512(_mm512_loadu_si512((const __m512i *)y[i].qs + l), _mm512_set1_epi8(-128)); + _mm512_storeu_si512((__m512i *)y[i].qs + l, v); + } +#else + _mm_storeu_si128((__m128i *)y[i].d, _mm256_cvtps_ph(_mm256_loadu_ps(dnew), _MM_ROUND_NEAREST)); +#endif + } + y += nb; + } +} + + +} // namespace + +bool iqk_set_kernels_kquants(int ne00, int typeA, int typeB, std::array& kernels, mul_mat_t& func16) { + + auto etypeA = ggml_type(typeA); + auto expected_type_B = etypeA == GGML_TYPE_IQ4_XS_R8 || etypeA == GGML_TYPE_Q4_K_R4 || etypeA == GGML_TYPE_Q5_K_R4 ? GGML_TYPE_Q8_K32 + //: etypeA == GGML_TYPE_Q8_K_R8 ? GGML_TYPE_Q8_KR8 + : etypeA == GGML_TYPE_Q8_KV || etypeA == GGML_TYPE_Q8_KV_R8 ? GGML_TYPE_Q8_KV + : etypeA == GGML_TYPE_Q4_K || etypeA == GGML_TYPE_Q5_K || + etypeA == GGML_TYPE_Q6_K ? GGML_TYPE_Q8_2_X4 + //etypeA == GGML_TYPE_Q6_K || etypeA == GGML_TYPE_Q3_K ? GGML_TYPE_Q8_2_X4 + //: etypeA == GGML_TYPE_Q4_K || etypeA == GGML_TYPE_Q5_K ? GGML_TYPE_Q8_2_X4 + : GGML_TYPE_Q8_K; + + if (ne00%QK_K != 0 || ggml_type(typeB) != expected_type_B) { + return false; + } + + func16 = nullptr; + + switch (typeA) { + case GGML_TYPE_Q2_K: + set_functions(kernels); + break; + case GGML_TYPE_Q3_K: + set_functions(kernels); + //IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qY_K_q8_2_X4_T, DequantizerQ3K_AVX2, kernels); + break; + case GGML_TYPE_Q4_K: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_2_X4_T, DequantizerQ4K_AVX2, kernels); + //set_functions(kernels); + break; + case GGML_TYPE_Q5_K: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_2_X4_T, DequantizerQ5K_AVX2, kernels); + //set_functions(kernels); + break; + case GGML_TYPE_Q6_K: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qY_K_q8_2_X4_T, DequantizerQ6K_AVX2, kernels); + //set_functions(kernels); + break; + case GGML_TYPE_IQ4_XS: + set_functions(kernels); + break; + case GGML_TYPE_Q2_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q2_k_r4_q8_k, kernels) + break; + case GGML_TYPE_Q3_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q3_k_r4_q8_k, kernels) + break; + case GGML_TYPE_Q4_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q4_k_r4_q8_k, kernels) + break; + case GGML_TYPE_Q5_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q5_k_r4_q8_k, kernels) + break; + case GGML_TYPE_Q6_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q6_k_r4_q8_k, kernels) + break; + case GGML_TYPE_IQ4_XS_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_iq4_xs_r8_q8_k_avx2, kernels) + break; + case GGML_TYPE_Q8_K_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_k_r8_q8_k, kernels) +//#ifdef HAVE_FANCY_SIMD +// func16 = mul_mat_q8_k_r8_q8_k<16>; +//#endif + break; +#ifdef HAVE_FANCY_SIMD + case GGML_TYPE_Q8_K_R16: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_k_r16_q8_k, kernels) + break; +#endif + case GGML_TYPE_Q8_KV: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_KV_q8_KV, kernels) +#ifdef HAVE_FANCY_SIMD + func16 = mul_mat_q8_KV_q8_KV<16>; +#endif + break; + case GGML_TYPE_Q8_KV_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_KV_r8_q8_KV, kernels); + break; + default: + return false; + } + + return true; + +} + +bool iqk_convert_kquants_q8X_r8(int type, int n, const void * vx, size_t bx, void * vy, int nrc_x) { + switch (ggml_type(type)) { + case GGML_TYPE_Q2_K: iqk_convert_q2_k_q8_k_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q3_K: iqk_convert_q3_k_q8_k_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q4_K: iqk_convert_q4_k_q8_1_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q5_K: iqk_convert_q5_k_q8_1_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q6_K: iqk_convert_q6_k_q8_0_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_IQ4_XS: iqk_convert_iq4_xs_q8_k_r8(n, vx, bx, vy, nrc_x); break; + default: return false; + } + return true; +} + +#else +// --------------------------------- __aarch64__ -------------------------------------- + +namespace { + +template +inline void accum_mins_8(const int16x8_t& mins, const Q8& q8, float32x4_t * acc, int i, float c) { + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + auto q8s = q8.load_bsums8(iy, i); + int32x4_t b1 = vmull_s16(vget_low_s16(mins), vget_low_s16(q8s)); + int32x4_t b2 = vmull_s16(vget_high_s16(mins), vget_high_s16(q8s)); + float32x4_t prod = vcvtq_f32_s32(vaddq_s32(b1, b2)); + acc[iy] = vmlaq_f32(acc[iy], prod, vdupq_n_f32(c*q8.scale(iy, i))); + } +} +template +inline void accum_mins_16(const int16x8x2_t& mins, const Q8& q8, float32x4_t * acc, int i, float c) { + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + auto q8s = q8.load_bsums(iy, i); + int32x4_t b1 = vmull_s16(vget_low_s16 (mins.val[0]), vget_low_s16 (q8s.val[0])); + int32x4_t b2 = vmull_s16(vget_high_s16(mins.val[0]), vget_high_s16(q8s.val[0])); + int32x4_t b3 = vmull_s16(vget_low_s16 (mins.val[1]), vget_low_s16 (q8s.val[1])); + int32x4_t b4 = vmull_s16(vget_high_s16(mins.val[1]), vget_high_s16(q8s.val[1])); + float32x4_t prod = vcvtq_f32_s32(vaddq_s32(vaddq_s32(b1, b2), vaddq_s32(b3, b4))); + acc[iy] = vmlaq_f32(acc[iy], prod, vdupq_n_f32(c*q8.scale(iy, i))); + } +} + +struct Scales8 { + uint32_t utmp[4]; + const uint8_t * sc8 = (const uint8_t *)utmp; + template + inline int32x4x2_t process_scales_mins(const Qx& x, const Q8& q8, int i, float32x4_t * acc) { + make_q4_scales(x.scales, utmp); + int16x8_t mins = vmovl_s8(vld1_s8((const int8_t *)sc8 + 8)); + accum_mins_8(mins, q8, acc, i, -GGML_FP16_TO_FP32(x.dmin)); + + uint8x8_t scales8 = vld1_u8(sc8); + uint16x8_t scales16 = vmovl_u8(scales8); + int32x4x2_t scales = {vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(scales16))), + vreinterpretq_s32_u32(vmovl_u16(vget_high_u16(scales16)))}; + return scales; + } + inline float32x4x4_t make_scales(float d, float m, const uint8_t * scales) { + make_q4_scales(scales, utmp); + auto d16 = vmovl_u8(vld1_u8(sc8+0)); + auto m16 = vmovl_u8(vld1_u8(sc8+8)); + auto vd = vdupq_n_f32(d); + auto vm = vdupq_n_f32(m); + return { vmulq_f32(vd, vcvtq_f32_u32(vmovl_u16(vget_low_u16 (d16)))), + vmulq_f32(vd, vcvtq_f32_u32(vmovl_u16(vget_high_u16(d16)))), + vmulq_f32(vm, vcvtq_f32_u32(vmovl_u16(vget_low_u16 (m16)))), + vmulq_f32(vm, vcvtq_f32_u32(vmovl_u16(vget_high_u16(m16)))) }; + } +}; + +struct DequantizerQ4K final : public BaseDequantizer { + DequantizerQ4K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} + + constexpr static int num_blocks() { return 8; } + constexpr static bool should_scale_quants() { return false; } + + template + inline int32x4x2_t new_block(int i, const Q8& q8, float32x4_t * acc) { + d = GGML_FP16_TO_FP32(x[i].d); + return s8.process_scales_mins(x[i], q8, i, acc); + } + inline float32x4x4_t new_block(int i) { + d = GGML_FP16_TO_FP32(x[i].d); + float m = -GGML_FP16_TO_FP32(x[i].dmin); + return s8.make_scales(d, m, x[i].scales); + } + inline void prepare(int i, int j) { + if (nrc == 1) bits.prepare_v2(x[i].qs+64*j); + else bits.prepare(x[i].qs+64*j); + } + + Q4bits bits; + Scales8 s8; + +}; + +struct HighBit5 { + const uint8x16_t mhb = vdupq_n_u8(0x10); + uint8x16x2_t bits; + inline void apply(uint8x16x4_t& b1, uint8x16x4_t& b2, bool do_shift) { + b1.val[0] = vorrq_u8(b1.val[0], vandq_u8(vshlq_n_u8(bits.val[0], 4), mhb)); + b1.val[1] = vorrq_u8(b1.val[1], vandq_u8(vshlq_n_u8(bits.val[1], 4), mhb)); + b1.val[2] = vorrq_u8(b1.val[2], vandq_u8(vshlq_n_u8(bits.val[0], 3), mhb)); + b1.val[3] = vorrq_u8(b1.val[3], vandq_u8(vshlq_n_u8(bits.val[1], 3), mhb)); + + b2.val[0] = vorrq_u8(b2.val[0], vandq_u8(vshlq_n_u8(bits.val[0], 2), mhb)); + b2.val[1] = vorrq_u8(b2.val[1], vandq_u8(vshlq_n_u8(bits.val[1], 2), mhb)); + b2.val[2] = vorrq_u8(b2.val[2], vandq_u8(vshlq_n_u8(bits.val[0], 1), mhb)); + b2.val[3] = vorrq_u8(b2.val[3], vandq_u8(vshlq_n_u8(bits.val[1], 1), mhb)); + + if (do_shift) { + bits.val[0] = vshrq_n_u8(bits.val[0], 4); + bits.val[1] = vshrq_n_u8(bits.val[1], 4); + } + } +}; + +struct HighBit3 { + const uint8x16_t mhb = vdupq_n_u8(0x04); + uint8x16x2_t bits; + inline void apply(uint8x16x4_t& b1, uint8x16x4_t& b2, bool do_shift) { + b1.val[0] = vorrq_u8(b1.val[0], vandq_u8(vshlq_n_u8(bits.val[0], 2), mhb)); + b1.val[1] = vorrq_u8(b1.val[1], vandq_u8(vshlq_n_u8(bits.val[1], 2), mhb)); + b1.val[2] = vorrq_u8(b1.val[2], vandq_u8(vshlq_n_u8(bits.val[0], 1), mhb)); + b1.val[3] = vorrq_u8(b1.val[3], vandq_u8(vshlq_n_u8(bits.val[1], 1), mhb)); + + b2.val[0] = vorrq_u8(b2.val[0], vandq_u8(bits.val[0], mhb)); + b2.val[1] = vorrq_u8(b2.val[1], vandq_u8(bits.val[1], mhb)); + b2.val[2] = vorrq_u8(b2.val[2], vandq_u8(vshrq_n_u8(bits.val[0], 1), mhb)); + b2.val[3] = vorrq_u8(b2.val[3], vandq_u8(vshrq_n_u8(bits.val[1], 1), mhb)); + + if (do_shift) { + bits.val[0] = vshrq_n_u8(bits.val[0], 4); + bits.val[1] = vshrq_n_u8(bits.val[1], 4); + } + } +}; + +struct DequantizerQ5K final : public BaseDequantizer { + DequantizerQ5K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} + + constexpr static int num_blocks() { return 8; } + constexpr static bool should_scale_quants() { return false; } + + template + inline int32x4x2_t new_block(int i, const Q8& q8, float32x4_t * acc) { + d = GGML_FP16_TO_FP32(x[i].d); + h.bits = vld1q_u8_x2(x[i].qh); + return s8.process_scales_mins(x[i], q8, i, acc); + } + inline float32x4x4_t new_block(int i) { + d = GGML_FP16_TO_FP32(x[i].d); + float m = -GGML_FP16_TO_FP32(x[i].dmin); + h.bits = vld1q_u8_x2(x[i].qh); + return s8.make_scales(d, m, x[i].scales); + } + inline void prepare(int i, int j) { + if (nrc == 1) bits.prepare_v2(x[i].qs+64*j); + else bits.prepare(x[i].qs+64*j); + h.apply(bits.b1, bits.b2, j == 0); + } + + Q4bits bits; + HighBit5 h; + Scales8 s8; + + uint8x16x2_t hbits; + +}; + +inline int32x4x4_t make_wider(const int16x8x2_t& scales16) { + int32x4x4_t scales = { + vmovl_s16(vget_low_s16 (scales16.val[0])), + vmovl_s16(vget_high_s16(scales16.val[0])), + vmovl_s16(vget_low_s16 (scales16.val[1])), + vmovl_s16(vget_high_s16(scales16.val[1])), + }; + return scales; +} + +template +inline int32x4x4_t process_scales_mins_16(const int8x16_t& scales8, const Q8& q8, float32x4_t * acc, int i, float c) { + int16x8x2_t scales16; + scales16.val[0] = vmovl_s8(vget_low_s8(scales8)); + scales16.val[1] = vmovl_s8(vget_high_s8(scales8)); + accum_mins_16(scales16, q8, acc, i, c); + return make_wider(scales16); +} + +struct DequantizerQ6K final : public BaseDequantizer { + DequantizerQ6K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} + + constexpr static int num_blocks() { return 16; } + constexpr static bool should_scale_quants() { return false; } + + template + inline int32x4x4_t new_block(int i, const Q8& q8, float32x4_t * acc) { + d = GGML_FP16_TO_FP32(x[i].d); + return process_scales_mins_16(vld1q_s8(x[i].scales), q8, acc, i, -32.f*d); + } + inline float32x4x4_t new_block(int i) { + d = GGML_FP16_TO_FP32(x[i].d); + auto vd = vdupq_n_f32(d); + auto scales8 = vld1q_s8(x[i].scales); + auto scales16_1 = vmovl_s8(vget_low_s8 (scales8)); + auto scales16_2 = vmovl_s8(vget_high_s8(scales8)); + return { vmulq_f32(vd, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (scales16_1)))), + vmulq_f32(vd, vcvtq_f32_s32(vmovl_s16(vget_high_s16(scales16_1)))), + vmulq_f32(vd, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (scales16_2)))), + vmulq_f32(vd, vcvtq_f32_s32(vmovl_s16(vget_high_s16(scales16_2)))) }; + } + + inline void prepare(int i, int j) { + + auto hbits = vld1q_u8_x2(x[i].qh + 32*j); + + bits.prepare64(x[i].ql+64*j); + bits.b1.val[0] = vorrq_u8(bits.b1.val[0], vandq_u8(vshlq_n_u8(hbits.val[0], 4), mhb)); + bits.b1.val[1] = vorrq_u8(bits.b1.val[1], vandq_u8(vshlq_n_u8(hbits.val[1], 4), mhb)); + bits.b1.val[2] = vorrq_u8(bits.b1.val[2], vandq_u8(vshlq_n_u8(hbits.val[0], 2), mhb)); + bits.b1.val[3] = vorrq_u8(bits.b1.val[3], vandq_u8(vshlq_n_u8(hbits.val[1], 2), mhb)); + + bits.b2.val[0] = vorrq_u8(bits.b2.val[0], vandq_u8(hbits.val[0], mhb)); + bits.b2.val[1] = vorrq_u8(bits.b2.val[1], vandq_u8(hbits.val[1], mhb)); + bits.b2.val[2] = vorrq_u8(bits.b2.val[2], vandq_u8(vshrq_n_u8(hbits.val[0], 2), mhb)); + bits.b2.val[3] = vorrq_u8(bits.b2.val[3], vandq_u8(vshrq_n_u8(hbits.val[1], 2), mhb)); + + } + + inline void prepare_signed(int i, int j) { + prepare(i, j); + auto m32 = vdupq_n_s8(-32); + for (int k = 0; k < 4; ++k) { + bits.b1.val[k] = vaddq_s8(bits.b1.val[k], m32); + bits.b2.val[k] = vaddq_s8(bits.b2.val[k], m32); + } + } + + Q4bits bits; + + const uint8x16_t mhb = vdupq_n_u8(0x30); + +}; + +struct DequantizerQ3K final : public BaseDequantizer { + DequantizerQ3K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} + + constexpr static int num_blocks() { return 16; } + constexpr static bool should_scale_quants() { return false; } + + template + inline int32x4x4_t new_block(int i, const Q8& q8, float32x4_t * acc) { + d = GGML_FP16_TO_FP32(x[i].d); + h.bits = vld1q_u8_x2(x[i].hmask); + mask = vdupq_n_u8(0x01); + const uint16_t * sc16 = (const uint16_t *)x[i].scales; + uint32_t aux0 = sc16[0] | (sc16[1] << 16); + uint32_t aux1 = sc16[2] | (sc16[3] << 16); + uint32_t aux2 = sc16[4] | (sc16[5] << 16); + aux32[0] = (aux0 & 0x0f0f0f0f) | ((aux2 << 4) & 0x30303030); + aux32[1] = (aux1 & 0x0f0f0f0f) | ((aux2 << 2) & 0x30303030); + aux32[2] = ((aux0 >> 4) & 0x0f0f0f0f) | ((aux2 >> 0) & 0x30303030); + aux32[3] = ((aux1 >> 4) & 0x0f0f0f0f) | ((aux2 >> 2) & 0x30303030); + auto scales8 = vaddq_s8(vld1q_s8((const int8_t *)aux32), vdupq_n_s8(-32)); + if (nrc > 1) { + return process_scales_mins_16(scales8, q8, acc, i, -4.f*d); + } + int16x8x2_t scales16; + scales16.val[0] = vmovl_s8(vget_low_s8(scales8)); + scales16.val[1] = vmovl_s8(vget_high_s8(scales8)); + return make_wider(scales16); + } + + inline void prepare(int i, int j) { + bits.prepare(x[i].qs+32*j); + if (nrc > 1) { + h.apply(bits.b1, bits.b2, j == 0); + } else { + auto minus4 = vdupq_n_u8(0xfc); + auto zero = vdupq_n_u8(0); + bits.b1.val[0] = vorrq_u8(bits.b1.val[0], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[0], mask), zero))); + bits.b1.val[1] = vorrq_u8(bits.b1.val[1], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[1], mask), zero))); + mask = vshlq_n_u8(mask, 1); + bits.b1.val[2] = vorrq_u8(bits.b1.val[2], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[0], mask), zero))); + bits.b1.val[3] = vorrq_u8(bits.b1.val[3], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[1], mask), zero))); + mask = vshlq_n_u8(mask, 1); + bits.b2.val[0] = vorrq_u8(bits.b2.val[0], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[0], mask), zero))); + bits.b2.val[1] = vorrq_u8(bits.b2.val[1], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[1], mask), zero))); + mask = vshlq_n_u8(mask, 1); + bits.b2.val[2] = vorrq_u8(bits.b2.val[2], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[0], mask), zero))); + bits.b2.val[3] = vorrq_u8(bits.b2.val[3], vandq_u8(minus4, vceqq_u8(vandq_u8(h.bits.val[1], mask), zero))); + mask = vshlq_n_u8(mask, 1); + } + } + + uint32_t aux32[4]; + + Q2bits bits; + + uint8x16_t mask; + HighBit3 h; + +}; + +struct DequantizerQ2K final : public BaseDequantizer { + DequantizerQ2K(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc) {} + + constexpr static int num_blocks() { return 16; } + constexpr static bool should_scale_quants() { return true; } + + template + inline void process_scales(int i, const Q8& q8, float32x4_t * acc) { + d = GGML_FP16_TO_FP32(x[i].d); + auto scales_and_mins = vld1q_u8(x[i].scales); + auto mins8 = vreinterpretq_s8_u8(vshrq_n_u8(scales_and_mins, 4)); + int16x8x2_t scales16; + scales16.val[0] = vmovl_s8(vget_low_s8(mins8)); + scales16.val[1] = vmovl_s8(vget_high_s8(mins8)); + accum_mins_16(scales16, q8, acc, i, -GGML_FP16_TO_FP32(x[i].dmin)); + + scales8 = vandq_u8(scales_and_mins, vdupq_n_u8(0xf)); + } + + template + inline int32x4x4_t new_block(int i, const Q8& q8, float32x4_t * acc) { + process_scales(i, q8, acc); + int16x8x2_t scales16; + scales16.val[0] = vmovl_s8(vget_low_s8(vreinterpretq_s8_u8(scales8))); + scales16.val[1] = vmovl_s8(vget_high_s8(vreinterpretq_s8_u8(scales8))); + return make_wider(scales16); + } + + template + inline void compute(const Q8& q8, int i, int j, int32x4_t * sumi) { + auto m1 = vdupq_n_u8(1); + auto shuffle = vdupq_n_u8(8*j); + bits.b1.val[0] = vmulq_u8(bits.b1.val[0], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); + bits.b1.val[1] = vmulq_u8(bits.b1.val[1], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); + bits.b1.val[2] = vmulq_u8(bits.b1.val[2], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); + bits.b1.val[3] = vmulq_u8(bits.b1.val[3], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); + bits.b2.val[0] = vmulq_u8(bits.b2.val[0], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); + bits.b2.val[1] = vmulq_u8(bits.b2.val[1], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); + bits.b2.val[2] = vmulq_u8(bits.b2.val[2], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); + bits.b2.val[3] = vmulq_u8(bits.b2.val[3], vqtbl1q_u8(scales8, shuffle)); shuffle = vaddq_u8(shuffle, m1); + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + auto q8b_1 = q8.load_quants(iy, i, 4*j+0); + sumi[iy] = ggml_vdotq_s32(ggml_vdotq_s32(sumi[iy], vreinterpretq_s8_u8(bits.b1.val[0]), q8b_1.val[0]), + vreinterpretq_s8_u8(bits.b1.val[1]), q8b_1.val[1]); + + auto q8b_2 = q8.load_quants(iy, i, 4*j+1); + sumi[iy] = ggml_vdotq_s32(ggml_vdotq_s32(sumi[iy], vreinterpretq_s8_u8(bits.b1.val[2]), q8b_2.val[0]), + vreinterpretq_s8_u8(bits.b1.val[3]), q8b_2.val[1]); + + auto q8b_3 = q8.load_quants(iy, i, 4*j+2); + sumi[iy] = ggml_vdotq_s32(ggml_vdotq_s32(sumi[iy], vreinterpretq_s8_u8(bits.b2.val[0]), q8b_3.val[0]), + vreinterpretq_s8_u8(bits.b2.val[1]), q8b_3.val[1]); + + auto q8b_4 = q8.load_quants(iy, i, 4*j+3); + sumi[iy] = ggml_vdotq_s32(ggml_vdotq_s32(sumi[iy], vreinterpretq_s8_u8(bits.b2.val[2]), q8b_4.val[0]), + vreinterpretq_s8_u8(bits.b2.val[3]), q8b_4.val[1]); + } + } + + inline void prepare(int i, int j) { + bits.prepare(x[i].qs+32*j); + } + + uint32_t aux32[4]; + + uint8x16_t scales8; + + Q2bits bits; + +}; + +struct DequantizerIQ4XS final : public BaseDequantizer { + + static int8x16_t load_values() { + static const int8_t iq4nl_values[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; + return vld1q_s8(iq4nl_values); + } + + DequantizerIQ4XS(const void * vx, size_t bx, int nrc) : BaseDequantizer(vx, bx, nrc), values(load_values()) {} + + constexpr static int num_blocks() { return 8; } + constexpr static bool should_scale_quants() { return false; } + + inline void new_row(int ix) { x = (const block_iq4_xs *)((const char *)vx + bx*ix); } + + template + inline int32x4x2_t new_block(int i, const Q8& q8, float32x4_t * acc) { + (void)q8; + (void)acc; + d = GGML_FP16_TO_FP32(x[i].d); + const uint16_t scales_h = x[i].scales_h; + const uint16_t * scales_l = (const uint16_t *)x[i].scales_l; + aux32[0] = scales_l[0] | (scales_l[1] << 16); + aux32[1] = aux32[0] >> 4; + // scl is ordered as 0, 2, 4, 6, 1, 3, 5, 7 + uint8x8_t scl8 = vand_u8(vld1_u8((const uint8_t *)aux32), vdup_n_u8(0xf)); + uint16_t * aux16 = (uint16_t *)aux32; + aux16[0] = scales_h << 4; aux16[1] = scales_h << 2; aux16[2] = scales_h; aux16[3] = scales_h >> 2; + // sch is ordered as 0, 4, 1, 5, 2, 6, 3, 7 + uint8x8_t sch8 = vand_u8(vld1_u8((const uint8_t *)aux16), vdup_n_u8(0x30)); + int8x8_t scales8 = vadd_s8(vreinterpret_s8_u8(vorr_u8(scl8, vtbl1_u8(sch8, vreinterpret_u8_u32(hshuff)))), vdup_n_s8(-32)); + // shuffle 0, 2, 4, 6, 1, 3, 5, 7 -> 0, 1, 2, 3, 4, 5, 6, 7 + scales8 = vtbl1_s8(scales8, vreinterpret_s8_u32(hshuff)); + int16x8_t scales16 = vmovl_s8(scales8); + int32x4x2_t scales = {vmovl_s16(vget_low_s16(scales16)), vmovl_s16(vget_high_s16(scales16))}; + return scales; + } + inline void prepare(int i, int j) { + bits.prepare16(x[i].qs+64*j); + //if (nrc == 1) { + // bits.prepare16_v2(x[i].qs+64*j); + //} else { + // bits.prepare16(x[i].qs+64*j); + //} + for (int k = 0; k < 4; ++k) { + bits.b1.val[k] = vreinterpretq_u8_s8(vqtbl1q_s8(values, bits.b1.val[k])); + bits.b2.val[k] = vreinterpretq_u8_s8(vqtbl1q_s8(values, bits.b2.val[k])); + } + } + + Q4bits bits; + const int8x16_t values; + uint32_t aux32[2]; + + constexpr static uint32x2_t hshuff = {0x05010400, 0x07030602}; + +}; + +IQK_ALWAYS_INLINE void prepare_q4_k_quants(const uint8x16_t& m4, const uint8x16x4_t& bits, int8x16_t * qx) { + qx[0] = vandq_u8(bits.val[0], m4); // 0...3 from the 4 rows + qx[1] = vandq_u8(bits.val[1], m4); // 16..19 + qx[2] = vandq_u8(bits.val[2], m4); // 4...7 + qx[3] = vandq_u8(bits.val[3], m4); // 20..23 + qx[4] = vshrq_n_u8(bits.val[0], 4); // 8..11 + qx[5] = vshrq_n_u8(bits.val[1], 4); // 24..27 + qx[6] = vshrq_n_u8(bits.val[2], 4); // 12..15 + qx[7] = vshrq_n_u8(bits.val[3], 4); // 28..31 +} + +template +void mul_mat_q2_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto mf = vdupq_n_u8(0x0f); + auto m03 = vdupq_n_u8(0x03); + int nbl = n / QK_K; + int8x16_t qx[4]; + float32x4_t acc[nrc_y] = {}; + int16x8x4_t i16scales; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q2_k_r4 * iq2 = (const block_q2_k_r4 *)((const char *)vx + ix*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { + int32x4_t isum[nrc_y] = {}; + auto d4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq2[ibl].d)); + auto m4 = vmulq_f32(vdupq_n_f32(-1.f), vcvt_f32_f16(vld1_f16((const float16_t *)iq2[ibl].d+4))); + for (int is = 0; is < 2; ++is) { + auto sl = vld1q_u8_x2(iq2[ibl].scales + 32*is); + auto m = vshrq_n_u8(sl.val[0], 4); + i16scales.val[0] = vmovl_u8(vget_low_u8 (m)); + i16scales.val[1] = vmovl_u8(vget_high_u8(m)); + m = vshrq_n_u8(sl.val[1], 4); + i16scales.val[2] = vmovl_u8(vget_low_u8 (m)); + i16scales.val[3] = vmovl_u8(vget_high_u8(m)); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = vdupq_n_s32(0); + auto bsums = vld1q_s16(q8.y[iy][ibl].bsums + 8*is); + auto b8 = vget_low_s16(bsums); + //auto bsums = q8.load_bsums(iy, ibl); + //auto b8 = vget_low_s16(bsums.val[0]); + sumi = vmlal_lane_s16(sumi, vget_low_s16 (i16scales.val[0]), b8, 0); + sumi = vmlal_lane_s16(sumi, vget_high_s16(i16scales.val[0]), b8, 1); + sumi = vmlal_lane_s16(sumi, vget_low_s16 (i16scales.val[1]), b8, 2); + sumi = vmlal_lane_s16(sumi, vget_high_s16(i16scales.val[1]), b8, 3); + b8 = vget_high_s16(bsums); + sumi = vmlal_lane_s16(sumi, vget_low_s16 (i16scales.val[2]), b8, 0); + sumi = vmlal_lane_s16(sumi, vget_high_s16(i16scales.val[2]), b8, 1); + sumi = vmlal_lane_s16(sumi, vget_low_s16 (i16scales.val[3]), b8, 2); + sumi = vmlal_lane_s16(sumi, vget_high_s16(i16scales.val[3]), b8, 3); + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(m4, vdupq_n_f32(q8.scale(iy, ibl))), vcvtq_f32_s32(sumi)); + } + m = vandq_u8(sl.val[0], mf); + i16scales.val[0] = vmovl_u8(vget_low_u8 (m)); + i16scales.val[1] = vmovl_u8(vget_high_u8(m)); + m = vandq_u8(sl.val[1], mf); + i16scales.val[2] = vmovl_u8(vget_low_u8 (m)); + i16scales.val[3] = vmovl_u8(vget_high_u8(m)); + for (int ib = 0; ib < 4; ++ib) { + auto bits = vld1q_u8_x2(iq2[ibl].qs + 128*is + 32*ib); + auto scales = vmovl_s16(vget_low_s16 (i16scales.val[ib])); + qx[0] = vreinterpretq_s8_u8(vandq_u8( bits.val[0], m03)); + qx[1] = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bits.val[0], 2), m03)); + qx[2] = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bits.val[0], 4), m03)); + qx[3] = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bits.val[0], 6), m03)); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8.y[iy][ibl].qs+128*is+32*ib); + auto sumi = interleaved_dotq(qx, y); + isum[iy] = vmlaq_s32(isum[iy], scales, sumi); + } + scales = vmovl_s16(vget_high_s16(i16scales.val[ib])); + qx[0] = vreinterpretq_s8_u8(vandq_u8( bits.val[1], m03)); + qx[1] = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bits.val[1], 2), m03)); + qx[2] = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bits.val[1], 4), m03)); + qx[3] = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(bits.val[1], 6), m03)); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8.y[iy][ibl].qs+128*is+32*ib+16); + auto sumi = interleaved_dotq(qx, y); + isum[iy] = vmlaq_s32(isum[iy], scales, sumi); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(d4, vdupq_n_f32(q8.scale(iy, ibl))), vcvtq_f32_s32(isum[iy])); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = vdupq_n_f32(0.f); + } + } +} + +template +void mul_mat_q3_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto mf = vdupq_n_u8(0x0f); + auto m30 = vdupq_n_u8(0x30); + auto m32 = vdupq_n_s8(-32); + auto m03 = vdupq_n_u8(0x03); + auto m04 = vdupq_n_u8(0x04); + int nbl = n / QK_K; + int8x16_t qx[4]; + float32x4_t acc[nrc_y] = {}; + int8x16x4_t i8scales; + int16x8x4_t i16scales; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q3_k_r4 * iq3 = (const block_q3_k_r4 *)((const char *)vx + ix*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { + int32x4_t isum[nrc_y] = {}; + auto d4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq3[ibl].d)); + auto sl = vld1q_u8_x2(iq3[ibl].scales_l); + auto sh = vld1q_u8(iq3[ibl].scales_h); + i8scales.val[0] = vaddq_s8(m32, vorrq_u8(vandq_u8(sl.val[0], mf), vandq_u8(vshlq_n_u8(sh, 4), m30))); + i8scales.val[1] = vaddq_s8(m32, vorrq_u8(vandq_u8(sl.val[1], mf), vandq_u8(vshlq_n_u8(sh, 2), m30))); + i8scales.val[2] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(sl.val[0], 4), vandq_u8(sh, m30))); + i8scales.val[3] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(sl.val[1], 4), vandq_u8(vshrq_n_u8(sh, 2), m30))); + for (int is = 0; is < 2; ++is) { + i16scales.val[0] = vmovl_s8(vget_low_s8 (i8scales.val[2*is+0])); + i16scales.val[1] = vmovl_s8(vget_high_s8(i8scales.val[2*is+0])); + i16scales.val[2] = vmovl_s8(vget_low_s8 (i8scales.val[2*is+1])); + i16scales.val[3] = vmovl_s8(vget_high_s8(i8scales.val[2*is+1])); + for (int ib = 0; ib < 4; ++ib) { + auto lbits = vld1q_u8_x2(iq3[ibl].qs + 128*is + 32*ib); + auto hbits = vld1q_u8(iq3[ibl].qh + 64*is + 16*ib); + hbits = veorq_u8(hbits, vdupq_n_u8(0xff)); + auto scales = vmovl_s16(vget_low_s16 (i16scales.val[ib])); + qx[0] = vsubq_s8(vreinterpretq_s8_u8(vandq_u8( lbits.val[0], m03)), vreinterpretq_s8_u8(vandq_u8(m04, vshlq_n_u8(hbits, 2)))); + qx[1] = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(lbits.val[0], 2), m03)), vreinterpretq_s8_u8(vandq_u8(m04, vshlq_n_u8(hbits, 1)))); + qx[2] = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(lbits.val[0], 4), m03)), vreinterpretq_s8_u8(vandq_u8(m04, hbits))); + qx[3] = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(lbits.val[0], 6), m03)), vreinterpretq_s8_u8(vandq_u8(m04, vshrq_n_u8(hbits, 1)))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8.y[iy][ibl].qs+128*is+32*ib); + auto sumi = interleaved_dotq(qx, y); + isum[iy] = vmlaq_s32(isum[iy], scales, sumi); + } + scales = vmovl_s16(vget_high_s16(i16scales.val[ib])); + qx[0] = vsubq_s8(vreinterpretq_s8_u8(vandq_u8( lbits.val[1], m03)), vreinterpretq_s8_u8(vandq_u8(m04, vshrq_n_u8(hbits, 2)))); + qx[1] = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(lbits.val[1], 2), m03)), vreinterpretq_s8_u8(vandq_u8(m04, vshrq_n_u8(hbits, 3)))); + qx[2] = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(lbits.val[1], 4), m03)), vreinterpretq_s8_u8(vandq_u8(m04, vshrq_n_u8(hbits, 4)))); + qx[3] = vsubq_s8(vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(lbits.val[1], 6), m03)), vreinterpretq_s8_u8(vandq_u8(m04, vshrq_n_u8(hbits, 5)))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8.y[iy][ibl].qs+128*is+32*ib+16); + auto sumi = interleaved_dotq(qx, y); + isum[iy] = vmlaq_s32(isum[iy], scales, sumi); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(d4, vdupq_n_f32(q8.scale(iy, ibl))), vcvtq_f32_s32(isum[iy])); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = vdupq_n_f32(0.f); + } + } +} + +template +void mul_mat_q4_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto mf = vdupq_n_u8(0xf); + auto m3 = vdupq_n_u8(0x30); + int nbl = n / QK_K; + int8x16_t qx[8]; + int8x16x2_t iscales; + int32x4x4_t scales; + float32x4_t acc[nrc_y] = {}; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q4_k_r4 * iq4 = (const block_q4_k_r4 *)((const char *)vx + ix*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { + auto d4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq4[ibl].d)); + auto m4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq4[ibl].d+4)); + m4 = vmulq_f32(m4, vdupq_n_f32(-1.f)); + auto sl = vld1q_u8_x2(iq4[ibl].scales_l); + auto sh = vld1q_u8(iq4[ibl].scales_h); + iscales.val[0] = vorrq_u8(vshrq_n_u8(sl.val[0], 4), vandq_u8(vshlq_n_u8(sh, 2), m3)); + iscales.val[1] = vorrq_u8(vshrq_n_u8(sl.val[1], 4), vandq_u8(vshrq_n_u8(sh, 2), m3)); + for (int is = 0; is < 2; ++is) { + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); + float32x4x4_t fscales; + fscales.val[0] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_1)))); + fscales.val[1] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_1)))); + fscales.val[2] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_2)))); + fscales.val[3] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_2)))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto m8 = vld1q_f32((const float *)q8.y[iy][ibl].bsums + 4*is); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[0], m8, 0); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[1], m8, 1); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[2], m8, 2); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[3], m8, 3); + } + } + iscales.val[0] = vorrq_u8(vandq_u8(sl.val[0], mf), vandq_u8(vshlq_n_u8(sh, 4), m3)); + iscales.val[1] = vorrq_u8(vandq_u8(sl.val[1], mf), vandq_u8(sh, m3)); + int32x4_t isum[nrc_y] = {}; + for (int is = 0; is < 2; ++is) { + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); + scales.val[0] = vmovl_s16(vget_low_s16(iscales16_1)); + scales.val[1] = vmovl_s16(vget_high_s16(iscales16_1)); + scales.val[2] = vmovl_s16(vget_low_s16(iscales16_2)); + scales.val[3] = vmovl_s16(vget_high_s16(iscales16_2)); + for (int ib = 0; ib < 4; ++ib) { + auto bits = vld1q_u8_x4(iq4[ibl].qs + 256*is + 64*ib); + prepare_q4_k_quants(mf, bits, qx); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][ibl].qs+128*is+32*ib); + auto sumi = interleaved_dotq(qx, y); + isum[iy] = vmlaq_s32(isum[iy], scales.val[ib], sumi); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(d4, vdupq_n_f32(q8.scale(iy, ibl))), vcvtq_f32_s32(isum[iy])); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = vdupq_n_f32(0.f); + } + } +} + +template +void mul_mat_q5_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto mf = vdupq_n_u8(0xf); + auto m30 = vdupq_n_u8(0x30); + auto m10 = vdupq_n_u8(0x10); + int nbl = n / QK_K; + int8x16_t qx[8]; + int8x16x2_t iscales; + int32x4x4_t scales; + float32x4_t acc[nrc_y] = {}; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q5_k_r4 * iq5 = (const block_q5_k_r4 *)((const char *)vx + ix*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { + auto d4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq5[ibl].d)); + auto m4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq5[ibl].d+4)); + m4 = vmulq_f32(m4, vdupq_n_f32(-1.f)); + auto sl = vld1q_u8_x2(iq5[ibl].scales_l); + auto sh = vld1q_u8(iq5[ibl].scales_h); + iscales.val[0] = vorrq_u8(vshrq_n_u8(sl.val[0], 4), vandq_u8(vshlq_n_u8(sh, 2), m30)); + iscales.val[1] = vorrq_u8(vshrq_n_u8(sl.val[1], 4), vandq_u8(vshrq_n_u8(sh, 2), m30)); + for (int is = 0; is < 2; ++is) { + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); + float32x4x4_t fscales; + fscales.val[0] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_1)))); + fscales.val[1] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_1)))); + fscales.val[2] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_low_s16(iscales16_2)))); + fscales.val[3] = vmulq_f32(m4, vcvtq_f32_s32(vmovl_s16(vget_high_s16(iscales16_2)))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto m8 = vld1q_f32((const float *)q8.y[iy][ibl].bsums + 4*is); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[0], m8, 0); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[1], m8, 1); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[2], m8, 2); + acc[iy] = vmlaq_laneq_f32(acc[iy], fscales.val[3], m8, 3); + } + } + iscales.val[0] = vorrq_u8(vandq_u8(sl.val[0], mf), vandq_u8(vshlq_n_u8(sh, 4), m30)); + iscales.val[1] = vorrq_u8(vandq_u8(sl.val[1], mf), vandq_u8(sh, m30)); + int32x4_t isum[nrc_y] = {}; + for (int is = 0; is < 2; ++is) { + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[is])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[is])); + scales.val[0] = vmovl_s16(vget_low_s16(iscales16_1)); + scales.val[1] = vmovl_s16(vget_high_s16(iscales16_1)); + scales.val[2] = vmovl_s16(vget_low_s16(iscales16_2)); + scales.val[3] = vmovl_s16(vget_high_s16(iscales16_2)); + for (int ib = 0; ib < 4; ++ib) { + auto lbits = vld1q_u8_x4(iq5[ibl].qs + 256*is + 64*ib); + auto hbits2 = vld1q_u8(iq5[ibl].qh + 64*is + 16*ib); + auto hbits1 = vshlq_n_u8(hbits2, 4); + prepare_q4_k_quants(mf, lbits, qx); + qx[0] = vorrq_u8(qx[0], vandq_u8(m10, hbits1)); + qx[1] = vorrq_u8(qx[1], vandq_u8(m10, hbits2)); + qx[2] = vorrq_u8(qx[2], vandq_u8(m10, vshrq_n_u8(hbits1, 2))); + qx[3] = vorrq_u8(qx[3], vandq_u8(m10, vshrq_n_u8(hbits2, 2))); + qx[4] = vorrq_u8(qx[4], vandq_u8(m10, vshrq_n_u8(hbits1, 1))); + qx[5] = vorrq_u8(qx[5], vandq_u8(m10, vshrq_n_u8(hbits2, 1))); + qx[6] = vorrq_u8(qx[6], vandq_u8(m10, vshrq_n_u8(hbits1, 3))); + qx[7] = vorrq_u8(qx[7], vandq_u8(m10, vshrq_n_u8(hbits2, 3))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][ibl].qs+128*is+32*ib); + auto sumi = interleaved_dotq(qx, y); + isum[iy] = vmlaq_s32(isum[iy], scales.val[ib], sumi); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(d4, vdupq_n_f32(q8.scale(iy, ibl))), vcvtq_f32_s32(isum[iy])); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = vdupq_n_f32(0.f); + } + } +} + +template +void mul_mat_q6_k_r4_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto mf = vdupq_n_u8(0x0f); + auto m3 = vdupq_n_u8(0x30); + auto m32 = vdupq_n_s8(-32); + int nbl = n / QK_K; + int8x16_t qx[4]; + float32x4_t acc[nrc_y] = {}; + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q6_k_r4 * iq6 = (const block_q6_k_r4 *)((const char *)vx + ix*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { + auto d4 = vcvt_f32_f16(vld1_f16((const float16_t *)iq6[ibl].d)); + int32x4_t isum[nrc_y] = {}; + for (int is = 0; is < 2; ++is) { + for (int ib = 0; ib < 4; ++ib) { + auto lbits = vld1q_u8_x4(iq6[ibl].ql + 256*is + 64*ib); + auto hbits = vld1q_u8(iq6[ibl].qh + 128*is + 32*ib); + auto iscales = vmovl_s8(vld1_s8(iq6[ibl].scales + 32*is + 8*ib)); + auto scales = vmovl_s16(vget_low_s16(iscales)); + qx[0] = vaddq_s8(m32, vorrq_u8(vandq_u8 (lbits.val[0], mf), vandq_u8(m3, vshlq_n_u8(hbits, 4)))); + qx[1] = vaddq_s8(m32, vorrq_u8(vandq_u8 (lbits.val[2], mf), vandq_u8(m3, hbits))); + qx[2] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(lbits.val[0], 4), vandq_u8(m3, vshlq_n_u8(hbits, 2)))); + qx[3] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(lbits.val[2], 4), vandq_u8(m3, vshrq_n_u8(hbits, 2)))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8.y[iy][ibl].qs+128*is+32*ib); + auto sumi = interleaved_dotq(qx, y); + isum[iy] = vmlaq_s32(isum[iy], scales, sumi); + } + scales = vmovl_s16(vget_high_s16(iscales)); + hbits = vld1q_u8(iq6[ibl].qh + 128*is + 32*ib + 16); + qx[0] = vaddq_s8(m32, vorrq_u8(vandq_u8 (lbits.val[1], mf), vandq_u8(m3, vshlq_n_u8(hbits, 4)))); + qx[1] = vaddq_s8(m32, vorrq_u8(vandq_u8 (lbits.val[3], mf), vandq_u8(m3, hbits))); + qx[2] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(lbits.val[1], 4), vandq_u8(m3, vshlq_n_u8(hbits, 2)))); + qx[3] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(lbits.val[3], 4), vandq_u8(m3, vshrq_n_u8(hbits, 2)))); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8.y[iy][ibl].qs+128*is+32*ib+16); + auto sumi = interleaved_dotq(qx, y); + isum[iy] = vmlaq_s32(isum[iy], scales, sumi); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(d4, vdupq_n_f32(q8.scale(iy, ibl))), vcvtq_f32_s32(isum[iy])); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = vdupq_n_f32(0.f); + } + } +} + +template +void mul_mat_q8_k_r8_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + int nbl = n / QK_K; + float32x4_t acc[2*nrc_y] = {}; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q8_k_r8 * iq8 = (const block_q8_k_r8 *)((const char *)vx + ix*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { + auto d4l = vcvt_f32_f16(vld1_f16((const float16_t *)iq8[ibl].d+0)); + auto d4h = vcvt_f32_f16(vld1_f16((const float16_t *)iq8[ibl].d+4)); + int32x4_t isum[2*nrc_y] = {}; + for (int ib = 0; ib < QK_K/16; ++ib) { + auto q1 = vld1q_s8_x4(iq8[ibl].qs + 128*ib + 0); + auto q2 = vld1q_s8_x4(iq8[ibl].qs + 128*ib + 64); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8.y[iy][ibl].qs+16*ib); + isum[2*iy+0] = vdotq_laneq_s32(isum[2*iy+0], q1.val[0], y, 0); + isum[2*iy+1] = vdotq_laneq_s32(isum[2*iy+1], q1.val[1], y, 0); + isum[2*iy+0] = vdotq_laneq_s32(isum[2*iy+0], q1.val[2], y, 1); + isum[2*iy+1] = vdotq_laneq_s32(isum[2*iy+1], q1.val[3], y, 1); + isum[2*iy+0] = vdotq_laneq_s32(isum[2*iy+0], q2.val[0], y, 2); + isum[2*iy+1] = vdotq_laneq_s32(isum[2*iy+1], q2.val[1], y, 2); + isum[2*iy+0] = vdotq_laneq_s32(isum[2*iy+0], q2.val[2], y, 3); + isum[2*iy+1] = vdotq_laneq_s32(isum[2*iy+1], q2.val[3], y, 3); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto d8 = vdupq_n_f32(q8.scale(iy, ibl)); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], vmulq_f32(d4l, d8), vcvtq_f32_s32(isum[2*iy+0])); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], vmulq_f32(d4h, d8), vcvtq_f32_s32(isum[2*iy+1])); + // Why did I have this? It is plain wrong! + //const float * bsum = (const float *)q8.y[iy][ibl].bsums; + //auto m8 = vdupq_n_f32(-128.f*bsum[0]); + //acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], d4l, m8); + //acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], d4l, m8); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix+0, iy, acc[2*iy+0]); + info.store(ix+4, iy, acc[2*iy+1]); + acc[2*iy+0] = acc[2*iy+1] = vdupq_n_f32(0.f); + } + } +} + +template +void mul_mat_iq4_xs_r8_q8_k(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto m4 = vdupq_n_u8(0xf); + auto m3 = vdupq_n_u8(0x30); + auto m32 = vdupq_n_s8(-32); + auto values = vld1q_s8(iq4k_values); + int nbl = n / QK_K; + int8x16_t qx[8]; + int8x16x4_t iscales; + int32x4x2_t scales; + float32x4_t acc[2*nrc_y] = {}; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_iq4_xs_r8 * iq4 = (const block_iq4_xs_r8 *)((const char *)vx + ix*bx); + for (int ibl = 0; ibl < nbl; ++ibl) { + auto d4_f16 = vld1q_f16((const float16_t *)iq4[ibl].d); + auto d4l = vcvt_f32_f16(vget_low_f16 (d4_f16)); + auto d4h = vcvt_f32_f16(vget_high_f16(d4_f16)); + auto sl = vld1q_u8_x2(iq4[ibl].scales_l); + auto sh = vld1q_u8(iq4[ibl].scales_h); + iscales.val[0] = vaddq_s8(vorrq_u8(vandq_u8(sl.val[0], m4), vandq_u8(vshlq_n_u8(sh, 4), m3)), m32); + iscales.val[1] = vaddq_s8(vorrq_u8(vandq_u8(sl.val[1], m4), vandq_u8(vshlq_n_u8(sh, 2), m3)), m32); + iscales.val[2] = vaddq_s8(vorrq_u8(vshrq_n_u8(sl.val[0], 4), vandq_u8(sh, m3)), m32); + iscales.val[3] = vaddq_s8(vorrq_u8(vshrq_n_u8(sl.val[1], 4), vandq_u8(vshrq_n_u8(sh, 2), m3)), m32); + int32x4_t isum[nrc_y] = {}; + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[ib64])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[ib64])); + scales.val[0] = vmovl_s16(vget_low_s16(iscales16_1)); + scales.val[1] = vmovl_s16(vget_low_s16(iscales16_2)); + for (int l = 0; l < 2; ++l) { + uint8x16x2_t bits; + bits.val[0] = vld1q_u8(iq4[ibl].qs + 256*ib64 + 128*l); + bits.val[1] = vld1q_u8(iq4[ibl].qs + 256*ib64 + 128*l + 32); + prepare_iq4_nl_quants_r8(values, m4, bits, qx+0); + bits.val[0] = vld1q_u8(iq4[ibl].qs + 256*ib64 + 128*l + 64); + bits.val[1] = vld1q_u8(iq4[ibl].qs + 256*ib64 + 128*l + 96); + prepare_iq4_nl_quants_r8(values, m4, bits, qx+4); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][ibl].qs+64*ib64+32*l); + auto sumi = vdupq_n_s32(0); + sumi = vdotq_laneq_s32(sumi, qx[0], y.val[0], 0); + sumi = vdotq_laneq_s32(sumi, qx[1], y.val[0], 1); + sumi = vdotq_laneq_s32(sumi, qx[2], y.val[0], 2); + sumi = vdotq_laneq_s32(sumi, qx[3], y.val[0], 3); + sumi = vdotq_laneq_s32(sumi, qx[4], y.val[1], 0); + sumi = vdotq_laneq_s32(sumi, qx[5], y.val[1], 1); + sumi = vdotq_laneq_s32(sumi, qx[6], y.val[1], 2); + sumi = vdotq_laneq_s32(sumi, qx[7], y.val[1], 3); + isum[iy] = vmlaq_s32(isum[iy], sumi, scales.val[l]); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto d8 = vdupq_n_f32(q8.scale(iy, ibl)); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], vmulq_f32(d4l, d8), vcvtq_f32_s32(isum[iy])); + isum[iy] = vdupq_n_s32(0); + } + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + auto iscales16_1 = vmovl_s8(vget_low_s8(iscales.val[ib64])); + auto iscales16_2 = vmovl_s8(vget_high_s8(iscales.val[ib64])); + scales.val[0] = vmovl_s16(vget_high_s16(iscales16_1)); + scales.val[1] = vmovl_s16(vget_high_s16(iscales16_2)); + for (int l = 0; l < 2; ++l) { + uint8x16x2_t bits; + bits.val[0] = vld1q_u8(iq4[ibl].qs + 256*ib64 + 128*l + 16); + bits.val[1] = vld1q_u8(iq4[ibl].qs + 256*ib64 + 128*l + 48); + prepare_iq4_nl_quants_r8(values, m4, bits, qx+0); + bits.val[0] = vld1q_u8(iq4[ibl].qs + 256*ib64 + 128*l + 80); + bits.val[1] = vld1q_u8(iq4[ibl].qs + 256*ib64 + 128*l +112); + prepare_iq4_nl_quants_r8(values, m4, bits, qx+4); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][ibl].qs+64*ib64+32*l); + auto sumi = vdupq_n_s32(0); + sumi = vdotq_laneq_s32(sumi, qx[0], y.val[0], 0); + sumi = vdotq_laneq_s32(sumi, qx[1], y.val[0], 1); + sumi = vdotq_laneq_s32(sumi, qx[2], y.val[0], 2); + sumi = vdotq_laneq_s32(sumi, qx[3], y.val[0], 3); + sumi = vdotq_laneq_s32(sumi, qx[4], y.val[1], 0); + sumi = vdotq_laneq_s32(sumi, qx[5], y.val[1], 1); + sumi = vdotq_laneq_s32(sumi, qx[6], y.val[1], 2); + sumi = vdotq_laneq_s32(sumi, qx[7], y.val[1], 3); + isum[iy] = vmlaq_s32(isum[iy], sumi, scales.val[l]); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto d8 = vdupq_n_f32(q8.scale(iy, ibl)); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], vmulq_f32(d4h, d8), vcvtq_f32_s32(isum[iy])); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix+0, iy, acc[2*iy+0]); + info.store(ix+4, iy, acc[2*iy+1]); + acc[2*iy+0] = acc[2*iy+1] = vdupq_n_f32(0.f); + } + } +} + +static void mul_mat_q8_KV_q8_KV_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(n%32 == 0); + int32x4_t acc[4] = {}; + auto dptr = (const float *)info.src1_row(0); + const float dy = dptr[0]; + auto q8y = (const int8_t *)(dptr + 2); + for (int ix = 0; ix < nrc_x; ++ix) { + auto dx = (const float *)((const char *)vx + ix*bx); + auto q8x = (const int8_t *)(dx + 2); + for (int i = 0; i < n/64; ++i) { + auto qx = vld1q_s8_x4(q8x + 64*i); + for (int j = 0; j < 4; ++j) { + acc[j] = ggml_vdotq_s32(acc[j], qx.val[j], vld1q_s8(q8y + 64*i + 16*j)); + } + } + if (int i = 2*(n/64); i < n/32) { + auto qx = vld1q_s8_x2(q8x + 32*i); + for (int j = 0; j < 2; ++j) { + acc[j] = ggml_vdotq_s32(acc[j], qx.val[j], vld1q_s8(q8y + 32*i + 16*j)); + } + } + acc[0] = vaddq_s32(acc[0], acc[1]); + acc[2] = vaddq_s32(acc[2], acc[3]); + acc[0] = vaddq_s32(acc[0], acc[2]); + info.store(ix, 0, dx[0]*dy*vaddvq_s32(acc[0])); + acc[0] = acc[1] = acc[2] = acc[3] = vdupq_n_s32(0); + } +} + +template +static void mul_mat_q8_KV_q8_KV(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + GGML_ASSERT(n%16 == 0); + int8x16_t qx[4]; + int32x4_t acc[nrc_y] = {}; + float dy[nrc_y]; + const int8_t * q8y[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) { + auto dptr = (const float *)info.src1_row(iy); + dy[iy] = dptr[0]; + q8y[iy] = (const int8_t *)(dptr + 2); + } + const int8_t * q8x[4]; + float dx[4]; + for (int ix = 0; ix < nrc_x; ix += 4) { + for (int kx = 0; kx < 4; ++kx) { + auto dptr = (const float *)((const char *)vx + (ix+kx)*bx); + dx[kx] = dptr[0]; + q8x[kx] = (const int8_t *)(dptr + 2); + } + for (int i = 0; i < n/16; ++i) { + for (int kx = 0; kx < 4; ++kx) qx[kx] = vld1q_s8(q8x[kx] + 16*i); + auto row01 = vtrnq_s32(qx[0], qx[1]); + auto row23 = vtrnq_s32(qx[2], qx[3]); + qx[0] = vtrn1q_s64(row01.val[0], row23.val[0]); + qx[1] = vtrn1q_s64(row01.val[1], row23.val[1]); + qx[2] = vtrn2q_s64(row01.val[0], row23.val[0]); + qx[3] = vtrn2q_s64(row01.val[1], row23.val[1]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8y[iy] + 16*i); + acc[iy] = vdotq_laneq_s32(acc[iy], qx[0], y, 0); + acc[iy] = vdotq_laneq_s32(acc[iy], qx[1], y, 1); + acc[iy] = vdotq_laneq_s32(acc[iy], qx[2], y, 2); + acc[iy] = vdotq_laneq_s32(acc[iy], qx[3], y, 3); + } + } + auto scales_x = vld1q_f32(dx); + for (int iy = 0; iy < nrc_y; ++iy) { + auto scale = vmulq_f32(scales_x, vdupq_n_f32(dy[iy])); + info.store(ix, iy, vmulq_f32(scale, vcvtq_f32_s32(acc[iy]))); + acc[iy] = vdupq_n_s32(0); + } + } +} + +template +void mul_mat_q8_KV_r8_q8_KV(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + int32x4_t acc[2*nrc_y] = {}; + float dy[nrc_y]; + const int8_t * q8y[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) { + auto dptr = (const float *)info.src1_row(iy); + dy[iy] = dptr[0]; + q8y[iy] = (const int8_t *)(dptr + 2); + } + for (int ix = 0; ix < nrc_x; ix += 8) { + const float * dptr = (const float *)((const char *)vx + ix*bx); + auto q8x = (const int8_t *)(dptr + 8); + for (int ib = 0; ib < n/16; ++ib) { + auto q1 = vld1q_s8_x4(q8x + 128*ib + 0); + auto q2 = vld1q_s8_x4(q8x + 128*ib + 64); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8(q8y[iy]+16*ib); + acc[2*iy+0] = vdotq_laneq_s32(acc[2*iy+0], q1.val[0], y, 0); + acc[2*iy+1] = vdotq_laneq_s32(acc[2*iy+1], q1.val[1], y, 0); + acc[2*iy+0] = vdotq_laneq_s32(acc[2*iy+0], q1.val[2], y, 1); + acc[2*iy+1] = vdotq_laneq_s32(acc[2*iy+1], q1.val[3], y, 1); + acc[2*iy+0] = vdotq_laneq_s32(acc[2*iy+0], q2.val[0], y, 2); + acc[2*iy+1] = vdotq_laneq_s32(acc[2*iy+1], q2.val[1], y, 2); + acc[2*iy+0] = vdotq_laneq_s32(acc[2*iy+0], q2.val[2], y, 3); + acc[2*iy+1] = vdotq_laneq_s32(acc[2*iy+1], q2.val[3], y, 3); + } + } + auto scale1_x = vld1q_f32(dptr+0); + auto scale2_x = vld1q_f32(dptr+4); + for (int iy = 0; iy < nrc_y; ++iy) { + auto scale_y = vdupq_n_f32(dy[iy]); + auto scale1 = vmulq_f32(scale1_x, scale_y); + auto scale2 = vmulq_f32(scale2_x, scale_y); + info.store(ix+0, iy, vmulq_f32(scale1, vcvtq_f32_s32(acc[2*iy+0]))); + info.store(ix+4, iy, vmulq_f32(scale2, vcvtq_f32_s32(acc[2*iy+1]))); + acc[2*iy+0] = acc[2*iy+1] = vdupq_n_s32(0.f); + } + } +} + +typedef struct { + ggml_half d[16]; + int8_t qs[8*QK8_1]; +} block_q8_1_r8; + +void iqk_convert_q2_k_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q2_K * x8[8]; + + block_q8_k_r8 * y = (block_q8_k_r8 *)vy; + + float32_t f_values[QK_K]; + uint32_t block[8]; + + int8x16x2_t xv[4]; + + auto ml = vdupq_n_u8(0x03); + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q2_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + auto vd = vdupq_n_f32(GGML_FP16_TO_FP32(x8[k][i].d)); + auto vm = vdupq_n_f32(-GGML_FP16_TO_FP32(x8[k][i].dmin)); + auto block_max = vdupq_n_f32(0); + for (int i128 = 0; i128 < 2; ++i128) { + auto bits = vld1q_u8_x2(x8[k][i].qs+32*i128); + xv[0].val[0] = vandq_u8(bits.val[0], ml); + xv[0].val[1] = vandq_u8(bits.val[1], ml); + xv[1].val[0] = vandq_u8(vshrq_n_u8(bits.val[0], 2), ml); + xv[1].val[1] = vandq_u8(vshrq_n_u8(bits.val[1], 2), ml); + xv[2].val[0] = vandq_u8(vshrq_n_u8(bits.val[0], 4), ml); + xv[2].val[1] = vandq_u8(vshrq_n_u8(bits.val[1], 4), ml); + xv[3].val[0] = vshrq_n_u8(bits.val[0], 6); + xv[3].val[1] = vshrq_n_u8(bits.val[1], 6); + for (int l = 0; l < 4; ++l) { + auto d1 = vdupq_n_s8(x8[k][i].scales[8*i128 + 2*l + 0] & 0xf); + auto d2 = vdupq_n_s8(x8[k][i].scales[8*i128 + 2*l + 1] & 0xf); + auto q1_8 = vmulq_s8(d1, xv[l].val[0]); + auto q2_8 = vmulq_s8(d2, xv[l].val[1]); + auto q1_16_1 = vmovl_s8(vget_low_s8 (q1_8)); + auto q1_16_2 = vmovl_s8(vget_high_s8(q1_8)); + auto q2_16_1 = vmovl_s8(vget_low_s8 (q2_8)); + auto q2_16_2 = vmovl_s8(vget_high_s8(q2_8)); + float32x4x4_t f1{vcvtq_f32_s32(vmovl_s16(vget_low_s16(q1_16_1))), vcvtq_f32_s32(vmovl_s16(vget_high_s16(q1_16_1))), + vcvtq_f32_s32(vmovl_s16(vget_low_s16(q1_16_2))), vcvtq_f32_s32(vmovl_s16(vget_high_s16(q1_16_2)))}; + float32x4x4_t f2{vcvtq_f32_s32(vmovl_s16(vget_low_s16(q2_16_1))), vcvtq_f32_s32(vmovl_s16(vget_high_s16(q2_16_1))), + vcvtq_f32_s32(vmovl_s16(vget_low_s16(q2_16_2))), vcvtq_f32_s32(vmovl_s16(vget_high_s16(q2_16_2)))}; + + auto m1 = vmulq_f32(vm, vcvtq_f32_s32(vdupq_n_s32(x8[k][i].scales[8*i128 + 2*l + 0] >> 4))); + auto m2 = vmulq_f32(vm, vcvtq_f32_s32(vdupq_n_s32(x8[k][i].scales[8*i128 + 2*l + 1] >> 4))); + + for (int j = 0; j < 4; ++j) { + f1.val[j] = vfmaq_f32(m1, vd, f1.val[j]); + f2.val[j] = vfmaq_f32(m2, vd, f2.val[j]); + } + vst1q_f32_x4(f_values + 128*i128 + 32*l + 0, f1); + vst1q_f32_x4(f_values + 128*i128 + 32*l + 16, f2); + + auto max1 = vmaxq_f32(vmaxq_f32(vabsq_f32(f1.val[0]), vabsq_f32(f1.val[1])), vmaxq_f32(vabsq_f32(f1.val[2]), vabsq_f32(f1.val[3]))); + auto max2 = vmaxq_f32(vmaxq_f32(vabsq_f32(f2.val[0]), vabsq_f32(f2.val[1])), vmaxq_f32(vabsq_f32(f2.val[2]), vabsq_f32(f2.val[3]))); + block_max = vmaxq_f32(block_max, vmaxq_f32(max1, max2)); + } + } + auto max = vmaxvq_f32(block_max); + float d = max / 127.f; + auto id = vdupq_n_f32(d != 0.0f ? 1/d : 0.0f); + y[i].d[k] = GGML_FP32_TO_FP16(d); + int16x8x4_t i16; + for (int ib32 = 0; ib32 < 8; ++ib32) { + auto v1 = vld1q_f32_x4(f_values + 32*ib32 + 0); + auto v2 = vld1q_f32_x4(f_values + 32*ib32 + 16); + i16.val[0] = vcombine_s16(vmovn_s32(vcvtnq_s32_f32(vmulq_f32(id, v1.val[0]))), vmovn_s32(vcvtnq_s32_f32(vmulq_f32(id, v1.val[1])))); + i16.val[1] = vcombine_s16(vmovn_s32(vcvtnq_s32_f32(vmulq_f32(id, v1.val[2]))), vmovn_s32(vcvtnq_s32_f32(vmulq_f32(id, v1.val[3])))); + i16.val[2] = vcombine_s16(vmovn_s32(vcvtnq_s32_f32(vmulq_f32(id, v2.val[0]))), vmovn_s32(vcvtnq_s32_f32(vmulq_f32(id, v2.val[1])))); + i16.val[3] = vcombine_s16(vmovn_s32(vcvtnq_s32_f32(vmulq_f32(id, v2.val[2]))), vmovn_s32(vcvtnq_s32_f32(vmulq_f32(id, v2.val[3])))); + vst1q_s8((int8_t *)block + 0, vcombine_s8(vmovn_s16(i16.val[0]), vmovn_s16(i16.val[1]))); + vst1q_s8((int8_t *)block + 16, vcombine_s8(vmovn_s16(i16.val[2]), vmovn_s16(i16.val[3]))); + auto q8 = (uint32_t *)y[i].qs + 64*ib32; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + } + } + } + y += nb; + } +} + +void iqk_convert_q3_k_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q3_K * x8[8]; + + block_q8_k_r8 * y = (block_q8_k_r8 *)vy; + + uint32_t block[8]; + int8x16x2_t xv[8]; + uint32_t aux32[4]; + + auto ml = vdupq_n_s8(0x03); + auto mh = vdupq_n_s8(0x04); + + union { int8x16_t vec; int8_t val[16]; } helper; + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q3_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + float d = GGML_FP16_TO_FP32(x8[k][i].d); + auto sc16 = (const uint16_t *)x8[k][i].scales; + uint32_t aux0 = sc16[0] | (sc16[1] << 16); + uint32_t aux1 = sc16[2] | (sc16[3] << 16); + uint32_t aux2 = sc16[4] | (sc16[5] << 16); + aux32[0] = (aux0 & 0x0f0f0f0f) | ((aux2 << 4) & 0x30303030); + aux32[1] = (aux1 & 0x0f0f0f0f) | ((aux2 << 2) & 0x30303030); + aux32[2] = ((aux0 >> 4) & 0x0f0f0f0f) | ((aux2 >> 0) & 0x30303030); + aux32[3] = ((aux1 >> 4) & 0x0f0f0f0f) | ((aux2 >> 2) & 0x30303030); + helper.vec = vaddq_s8(vld1q_s8((const int8_t *)aux32), vdupq_n_s8(-32)); + auto hbits = vld1q_u8_x2(x8[k][i].hmask); + auto max_i16 = vdupq_n_u16(0); + for (int i128 = 0; i128 < 2; ++i128) { + auto q2bits = vld1q_u8_x2(x8[k][i].qs + 32*i128); + xv[4*i128+0].val[0] = vsubq_s8(vorrq_s8(vandq_s8(q2bits.val[0], ml), vandq_s8(vshlq_n_u8(hbits.val[0], 2), mh)), mh); + xv[4*i128+0].val[1] = vsubq_s8(vorrq_s8(vandq_s8(q2bits.val[1], ml), vandq_s8(vshlq_n_u8(hbits.val[1], 2), mh)), mh); + xv[4*i128+1].val[0] = vsubq_s8(vorrq_s8(vandq_s8(vshrq_n_u8(q2bits.val[0], 2), ml), vandq_s8(vshlq_n_u8(hbits.val[0], 1), mh)), mh); + xv[4*i128+1].val[1] = vsubq_s8(vorrq_s8(vandq_s8(vshrq_n_u8(q2bits.val[1], 2), ml), vandq_s8(vshlq_n_u8(hbits.val[1], 1), mh)), mh); + xv[4*i128+2].val[0] = vsubq_s8(vorrq_s8(vandq_s8(vshrq_n_u8(q2bits.val[0], 4), ml), vandq_s8(hbits.val[0], mh)), mh); + xv[4*i128+2].val[1] = vsubq_s8(vorrq_s8(vandq_s8(vshrq_n_u8(q2bits.val[1], 4), ml), vandq_s8(hbits.val[1], mh)), mh); + xv[4*i128+3].val[0] = vsubq_s8(vorrq_s8(vshrq_n_u8(q2bits.val[0], 6), vandq_s8(vshrq_n_u8(hbits.val[0], 1), mh)), mh); + xv[4*i128+3].val[1] = vsubq_s8(vorrq_s8(vshrq_n_u8(q2bits.val[1], 6), vandq_s8(vshrq_n_u8(hbits.val[1], 1), mh)), mh); + hbits.val[0] = vshrq_n_u8(hbits.val[0], 4); + hbits.val[1] = vshrq_n_u8(hbits.val[1], 4); + + for (int l = 0; l < 4; ++l) { + auto s1 = vdup_n_s8(helper.val[8*i128+2*l+0]); + auto s2 = vdup_n_s8(helper.val[8*i128+2*l+1]); + auto q16_1 = vmull_s8(s1, vget_low_s8 (xv[4*i128+l].val[0])); + auto q16_2 = vmull_s8(s1, vget_high_s8(xv[4*i128+l].val[0])); + auto q16_3 = vmull_s8(s2, vget_low_s8 (xv[4*i128+l].val[1])); + auto q16_4 = vmull_s8(s2, vget_high_s8(xv[4*i128+l].val[1])); + auto max1 = vmaxq_s16(vabsq_s16(q16_1), vabsq_s16(q16_2)); + auto max2 = vmaxq_s16(vabsq_s16(q16_3), vabsq_s16(q16_4)); + max_i16 = vmaxq_s16(max_i16, vmaxq_s16(max1, max2)); + } + } + auto imax16 = vmaxvq_s16(max_i16); + bool needs_scaling = true; + float dnew = float(imax16) / 127; + if (dnew < 1.f) { + dnew = 1.f; needs_scaling = false; + } + d *= dnew; + y[i].d[k] = GGML_FP32_TO_FP16(d); + auto scale = vdupq_n_f32(std::abs(dnew) > 1e-9f ? 1/dnew : 0.f); + for (int ib32 = 0; ib32 < 8; ++ib32) { + auto s1 = vdup_n_s8(helper.val[2*ib32+0]); + auto s2 = vdup_n_s8(helper.val[2*ib32+1]); + auto q16_1 = vmull_s8(s1, vget_low_s8 (xv[ib32].val[0])); + auto q16_2 = vmull_s8(s1, vget_high_s8(xv[ib32].val[0])); + auto q16_3 = vmull_s8(s2, vget_low_s8 (xv[ib32].val[1])); + auto q16_4 = vmull_s8(s2, vget_high_s8(xv[ib32].val[1])); + if (needs_scaling) { + int32x4x4_t i1{vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (q16_1))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(q16_1))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (q16_2))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(q16_2)))))}; + int32x4x4_t i2{vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (q16_3))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(q16_3))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (q16_4))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(q16_4)))))}; + int16x8x4_t i3{vcombine_s16(vmovn_s32(i1.val[0]), vmovn_s32(i1.val[1])), + vcombine_s16(vmovn_s32(i1.val[2]), vmovn_s32(i1.val[3])), + vcombine_s16(vmovn_s32(i2.val[0]), vmovn_s32(i2.val[1])), + vcombine_s16(vmovn_s32(i2.val[2]), vmovn_s32(i2.val[3]))}; + vst1q_s8((int8_t *)block + 0, vcombine_s8(vmovn_s16(i3.val[0]), vmovn_s16(i3.val[1]))); + vst1q_s8((int8_t *)block + 16, vcombine_s8(vmovn_s16(i3.val[2]), vmovn_s16(i3.val[3]))); + } else { + vst1q_s8((int8_t *)block + 0, vcombine_s8(vmovn_s16(q16_1), vmovn_s16(q16_2))); + vst1q_s8((int8_t *)block + 16, vcombine_s8(vmovn_s16(q16_3), vmovn_s16(q16_4))); + } + auto qs = (uint32_t *)y[i].qs + 64*ib32; + for (int l = 0; l < 8; ++l) { + qs[8*l + k] = block[l]; + } + } + } + } + y += nb; + } +} + +void iqk_convert_q6_k_q8_0_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q6_K * x8[8]; + + block_q8_0_r8 * y = (block_q8_0_r8 *)vy; + + float all_s[64]; + uint32_t block[8]; + int8x16x2_t xv[8]; + + auto ml = vdupq_n_u8(0x0f); + auto mh = vdupq_n_u8(0x30); + auto m32 = vdupq_n_s8(-32); + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q6_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + float d = GGML_FP16_TO_FP32(x8[k][i].d); + auto ql = x8[k][i].ql; + auto qh = x8[k][i].qh; + for (int i128 = 0; i128 < 2; ++i128) { + auto lbits1 = vld1q_u8_x2(ql + 64*i128 + 0); + auto lbits2 = vld1q_u8_x2(ql + 64*i128 + 32); + auto hbits = vld1q_u8_x2(qh + 32*i128); + xv[4*i128+0].val[0] = vaddq_s8(m32, vorrq_u8(vandq_u8(lbits1.val[0], ml), vandq_u8(vshlq_n_u8(hbits.val[0], 4), mh))); + xv[4*i128+0].val[1] = vaddq_s8(m32, vorrq_u8(vandq_u8(lbits1.val[1], ml), vandq_u8(vshlq_n_u8(hbits.val[1], 4), mh))); + xv[4*i128+1].val[0] = vaddq_s8(m32, vorrq_u8(vandq_u8(lbits2.val[0], ml), vandq_u8(vshlq_n_u8(hbits.val[0], 2), mh))); + xv[4*i128+1].val[1] = vaddq_s8(m32, vorrq_u8(vandq_u8(lbits2.val[1], ml), vandq_u8(vshlq_n_u8(hbits.val[1], 2), mh))); + xv[4*i128+2].val[0] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(lbits1.val[0], 4), vandq_u8(hbits.val[0], mh))); + xv[4*i128+2].val[1] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(lbits1.val[1], 4), vandq_u8(hbits.val[1], mh))); + xv[4*i128+3].val[0] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(lbits2.val[0], 4), vandq_u8(vshrq_n_u8(hbits.val[0], 2), mh))); + xv[4*i128+3].val[1] = vaddq_s8(m32, vorrq_u8(vshrq_n_u8(lbits2.val[1], 4), vandq_u8(vshrq_n_u8(hbits.val[1], 2), mh))); + } + for (int ib32 = 0; ib32 < 8; ++ib32) { + // We have two blocks of 16 with different scales + // We multiply the quants with the scales, find the max value, and convert to 8-bit quants with a single block scale. + auto s1 = vdup_n_s8(x8[k][i].scales[2*ib32+0]); + auto s2 = vdup_n_s8(x8[k][i].scales[2*ib32+1]); + int16x8x4_t i16{vmull_s8(s1, vget_low_s8(xv[ib32].val[0])), vmull_s8(s1, vget_high_s8(xv[ib32].val[0])), + vmull_s8(s2, vget_low_s8(xv[ib32].val[1])), vmull_s8(s2, vget_high_s8(xv[ib32].val[1]))}; + auto imax16 = vmaxq_u16(vmaxq_u16(vabsq_s16(i16.val[0]), vabsq_s16(i16.val[1])), vmaxq_u16(vabsq_s16(i16.val[2]), vabsq_s16(i16.val[3]))); + auto imax = vmaxvq_u16(imax16); + float max = float(imax) / 127; + all_s[8*ib32+k] = d*max; + if (max > 1e-9f) { + auto scale = vdupq_n_f32(1/max); + int32x4x4_t i32_1 = {vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (i16.val[0]))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(i16.val[0]))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (i16.val[1]))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(i16.val[1])))))}; + int32x4x4_t i32_2 = {vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (i16.val[2]))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(i16.val[2]))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (i16.val[3]))))), + vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(i16.val[3])))))}; + i16.val[0] = vcombine_s16(vmovn_s32(i32_1.val[0]), vmovn_s32(i32_1.val[1])); + i16.val[1] = vcombine_s16(vmovn_s32(i32_1.val[2]), vmovn_s32(i32_1.val[3])); + i16.val[2] = vcombine_s16(vmovn_s32(i32_2.val[0]), vmovn_s32(i32_2.val[1])); + i16.val[3] = vcombine_s16(vmovn_s32(i32_2.val[2]), vmovn_s32(i32_2.val[3])); + vst1q_s8((int8_t *)block + 0, vcombine_s8(vmovn_s16(i16.val[0]), vmovn_s16(i16.val[1]))); + vst1q_s8((int8_t *)block + 16, vcombine_s8(vmovn_s16(i16.val[2]), vmovn_s16(i16.val[3]))); + } else { + std::memset(block, 0, 8*sizeof(uint32_t)); + } + auto qs = (uint32_t *)y[ib32].qs; + for (int l = 0; l < 4; ++l) { + qs[8*l + k + 0] = block[l + 0]; + qs[8*l + k + 32] = block[l + 4]; + } + } + } + for (int ib32 = 0; ib32 < 8; ++ib32) { + vst1_f16((float16_t *)y[ib32].d + 0, vcvt_f16_f32(vld1q_f32(all_s + 8*ib32 + 0))); + vst1_f16((float16_t *)y[ib32].d + 4, vcvt_f16_f32(vld1q_f32(all_s + 8*ib32 + 4))); + } + y += QK_K/32; + } + } +} + +template +static void mul_mat_q6_k_q8_0_x4(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + Q8 q8(info); + + DequantizerQ6K deq(vx, bx, nrc_y); + + for (int ix = 0; ix < nrc_x; ++ix) { + + deq.new_row(ix); + + float32x4_t acc[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) acc[iy] = vdupq_n_f32(0.f); + + for (int i = 0; i < nb; ++i) { + + auto scales = deq.new_block(i); + + deq.prepare_signed(i, 0); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][2*i+0].qs); + auto dot1 = vpaddq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[0], y.val[0]), ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[1], y.val[1])); + y = vld1q_s8_x2(q8.y[iy][2*i+0].qs+32); + auto dot2 = vpaddq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[2], y.val[0]), ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[3], y.val[1])); + auto dot12 = vpaddq_s32(dot1, dot2); // 0, 1, 2, 3 + y = vld1q_s8_x2(q8.y[iy][2*i+0].qs+64); + auto dot3 = vpaddq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[0], y.val[0]), ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[1], y.val[1])); + y = vld1q_s8_x2(q8.y[iy][2*i+0].qs+96); + auto dot4 = vpaddq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[2], y.val[0]), ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[3], y.val[1])); + auto dot34 = vpaddq_s32(dot3, dot4); // 4, 5, 6, 7 + auto d8 = vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][2*i+0].d)); + auto d8_1 = vzip1q_f32(d8, d8); + auto d8_2 = vzip2q_f32(d8, d8); + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(scales.val[0], d8_1), vcvtq_f32_s32(dot12)); + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(scales.val[1], d8_2), vcvtq_f32_s32(dot34)); + } + + deq.prepare_signed(i, 1); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][2*i+1].qs); + auto dot1 = vpaddq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[0], y.val[0]), ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[1], y.val[1])); + y = vld1q_s8_x2(q8.y[iy][2*i+1].qs+32); + auto dot2 = vpaddq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[2], y.val[0]), ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[3], y.val[1])); + auto dot12 = vpaddq_s32(dot1, dot2); // 0, 1, 2, 3 + y = vld1q_s8_x2(q8.y[iy][2*i+1].qs+64); + auto dot3 = vpaddq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[0], y.val[0]), ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[1], y.val[1])); + y = vld1q_s8_x2(q8.y[iy][2*i+1].qs+96); + auto dot4 = vpaddq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[2], y.val[0]), ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[3], y.val[1])); + auto dot34 = vpaddq_s32(dot3, dot4); // 4, 5, 6, 7 + auto d8 = vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][2*i+1].d)); + auto d8_1 = vzip1q_f32(d8, d8); + auto d8_2 = vzip2q_f32(d8, d8); + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(scales.val[2], d8_1), vcvtq_f32_s32(dot12)); + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(scales.val[3], d8_2), vcvtq_f32_s32(dot34)); + } + } + + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, vaddvq_f32(acc[iy])); + } + } +} + +void iqk_convert_q4_k_q8_1_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q4_K * x8[8]; + + block_q8_1_r8 * y = (block_q8_1_r8 *)vy; + + ggml_half dh[16]; + uint16_t all_ls[128]; + + uint32_t utmp[4]; + const uint8_t * u8 = (const uint8_t *)utmp; + uint32_t block[8]; + + auto ml = vdupq_n_u8(0xf); + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q4_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + dh[k+0] = x8[k][i].d; + dh[k+8] = x8[k][i].dmin; + make_q4_scales(x8[k][i].scales, utmp); + auto qs = x8[k][i].qs; + for (int ib64 = 0; ib64 < 4; ++ib64) { + all_ls[8*(2*ib64 + 0) + k ] = u8[2*ib64+0]; + all_ls[8*(2*ib64 + 1) + k ] = u8[2*ib64+1]; + all_ls[8*(2*ib64 + 0) + k + 64] = u8[2*ib64+8]; + all_ls[8*(2*ib64 + 1) + k + 64] = u8[2*ib64+9]; + auto bits = vld1q_u8_x2(qs+32*ib64); + uint8x16x2_t xv1{vandq_u8(bits.val[0], ml), vandq_u8(bits.val[1], ml)}; + uint8x16x2_t xv2{vshrq_n_u8(bits.val[0], 4), vshrq_n_u8(bits.val[1], 4)}; + vst1q_u8_x2((uint8_t *)block, xv1); + auto q8 = (uint32_t *)y[2*ib64+0].qs; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + vst1q_u8_x2((uint8_t *)block, xv2); + q8 = (uint32_t *)y[2*ib64+1].qs; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + } + } + float32x4x2_t vd{ vcvt_f32_f16(vld1_f16((const float16_t *)dh+0)), vcvt_f32_f16(vld1_f16((const float16_t *)dh+ 4)) }; + float32x4x2_t vm{ vcvt_f32_f16(vld1_f16((const float16_t *)dh+8)), vcvt_f32_f16(vld1_f16((const float16_t *)dh+12)) }; + vm.val[0] = vmulq_f32(vdupq_n_f32(-1.f), vm.val[0]); + vm.val[1] = vmulq_f32(vdupq_n_f32(-1.f), vm.val[1]); + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + auto iscales16 = vld1q_u16(all_ls + 8*ib32); + uint32x4x2_t iscales32 = { vmovl_u16(vget_low_u16(iscales16)), vmovl_u16(vget_high_u16(iscales16)) }; + auto scales1 = vmulq_f32(vd.val[0], vcvtq_f32_u32(iscales32.val[0])); + auto scales2 = vmulq_f32(vd.val[1], vcvtq_f32_u32(iscales32.val[1])); + vst1_f16((float16_t *)y[ib32].d+0, vcvt_f16_f32(scales1)); + vst1_f16((float16_t *)y[ib32].d+4, vcvt_f16_f32(scales2)); + + iscales16 = vld1q_u16(all_ls + 8*ib32 + 64); + iscales32 = { vmovl_u16(vget_low_u16(iscales16)), vmovl_u16(vget_high_u16(iscales16)) }; + scales1 = vmulq_f32(vm.val[0], vcvtq_f32_u32(iscales32.val[0])); + scales2 = vmulq_f32(vm.val[1], vcvtq_f32_u32(iscales32.val[1])); + vst1_f16((float16_t *)y[ib32].d+ 8, vcvt_f16_f32(scales1)); + vst1_f16((float16_t *)y[ib32].d+12, vcvt_f16_f32(scales2)); + } + y += QK_K/32; + } + } +} + +void iqk_convert_q5_k_q8_1_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_q5_K * x8[8]; + + block_q8_1_r8 * y = (block_q8_1_r8 *)vy; + + ggml_half dh[16]; + uint16_t all_ls[128]; + + uint32_t utmp[4]; + const uint8_t * u8 = (const uint8_t *)utmp; + uint32_t block[8]; + + auto ml = vdupq_n_u8(0x0f); + auto mh = vdupq_n_u8(0x10); + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_q5_K *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + dh[k+0] = x8[k][i].d; + dh[k+8] = x8[k][i].dmin; + make_q4_scales(x8[k][i].scales, utmp); + auto qs = x8[k][i].qs; + auto hbits = vld1q_u8_x2(x8[k][i].qh); + for (int ib64 = 0; ib64 < 4; ++ib64) { + all_ls[8*(2*ib64 + 0) + k ] = u8[2*ib64+0]; + all_ls[8*(2*ib64 + 1) + k ] = u8[2*ib64+1]; + all_ls[8*(2*ib64 + 0) + k + 64] = u8[2*ib64+8]; + all_ls[8*(2*ib64 + 1) + k + 64] = u8[2*ib64+9]; + auto bits = vld1q_u8_x2(qs+32*ib64); + uint8x16x2_t xv1{vandq_u8(bits.val[0], ml), vandq_u8(bits.val[1], ml)}; + uint8x16x2_t xv2{vshrq_n_u8(bits.val[0], 4), vshrq_n_u8(bits.val[1], 4)}; + xv1.val[0] = vorrq_u8(xv1.val[0], vandq_u8(vshlq_n_u8(hbits.val[0], 4), mh)); + xv1.val[1] = vorrq_u8(xv1.val[1], vandq_u8(vshlq_n_u8(hbits.val[1], 4), mh)); + xv2.val[0] = vorrq_u8(xv2.val[0], vandq_u8(vshlq_n_u8(hbits.val[0], 3), mh)); + xv2.val[1] = vorrq_u8(xv2.val[1], vandq_u8(vshlq_n_u8(hbits.val[1], 3), mh)); + vst1q_u8_x2((uint8_t *)block, xv1); + auto q8 = (uint32_t *)y[2*ib64+0].qs; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + vst1q_u8_x2((uint8_t *)block, xv2); + q8 = (uint32_t *)y[2*ib64+1].qs; + for (int l = 0; l < 4; ++l) { + q8[8*l + k + 0] = block[l + 0]; + q8[8*l + k + 32] = block[l + 4]; + } + hbits.val[0] = vshrq_n_u8(hbits.val[0], 2); + hbits.val[1] = vshrq_n_u8(hbits.val[1], 2); + } + } + float32x4x2_t vd{ vcvt_f32_f16(vld1_f16((const float16_t *)dh+0)), vcvt_f32_f16(vld1_f16((const float16_t *)dh+ 4)) }; + float32x4x2_t vm{ vcvt_f32_f16(vld1_f16((const float16_t *)dh+8)), vcvt_f32_f16(vld1_f16((const float16_t *)dh+12)) }; + vm.val[0] = vmulq_f32(vdupq_n_f32(-1.f), vm.val[0]); + vm.val[1] = vmulq_f32(vdupq_n_f32(-1.f), vm.val[1]); + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + auto iscales16 = vld1q_u16(all_ls + 8*ib32); + uint32x4x2_t iscales32 = { vmovl_u16(vget_low_u16(iscales16)), vmovl_u16(vget_high_u16(iscales16)) }; + auto scales1 = vmulq_f32(vd.val[0], vcvtq_f32_u32(iscales32.val[0])); + auto scales2 = vmulq_f32(vd.val[1], vcvtq_f32_u32(iscales32.val[1])); + vst1_f16((float16_t *)y[ib32].d+0, vcvt_f16_f32(scales1)); + vst1_f16((float16_t *)y[ib32].d+4, vcvt_f16_f32(scales2)); + + iscales16 = vld1q_u16(all_ls + 8*ib32 + 64); + iscales32 = { vmovl_u16(vget_low_u16(iscales16)), vmovl_u16(vget_high_u16(iscales16)) }; + scales1 = vmulq_f32(vm.val[0], vcvtq_f32_u32(iscales32.val[0])); + scales2 = vmulq_f32(vm.val[1], vcvtq_f32_u32(iscales32.val[1])); + vst1_f16((float16_t *)y[ib32].d+ 8, vcvt_f16_f32(scales1)); + vst1_f16((float16_t *)y[ib32].d+12, vcvt_f16_f32(scales2)); + } + y += QK_K/32; + } + } +} + +template +static void mul_mat_qX_k_q8_1_x4(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n % QK_K == 0); + const int nb = n / QK_K; + + Q8 q8(info); + + Dequantizer deq(vx, bx, nrc_y); + + for (int ix = 0; ix < nrc_x; ++ix) { + + deq.new_row(ix); + + float32x4_t acc[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) acc[iy] = vdupq_n_f32(0.f); + + for (int i = 0; i < nb; ++i) { + + auto scales = deq.new_block(i); + + for (int iy = 0; iy < nrc_y; ++iy) { + auto m1 = vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][2*i+0].d+4)); + auto m2 = vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][2*i+1].d+4)); + acc[iy] = vfmaq_f32(acc[iy], scales.val[2], m1); + acc[iy] = vfmaq_f32(acc[iy], scales.val[3], m2); + } + + deq.prepare(i, 0); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][2*i+0].qs); + auto dot1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[0], y.val[0]), deq.bits.b1.val[1], y.val[1]); + y = vld1q_s8_x2(q8.y[iy][2*i+0].qs+32); + auto dot2 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[2], y.val[0]), deq.bits.b1.val[3], y.val[1]); + auto dot12 = vpaddq_s32(dot1, dot2); // 0, 0, 1, 1 + y = vld1q_s8_x2(q8.y[iy][2*i+0].qs+64); + auto dot3 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[0], y.val[0]), deq.bits.b2.val[1], y.val[1]); + y = vld1q_s8_x2(q8.y[iy][2*i+0].qs+96); + auto dot4 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[2], y.val[0]), deq.bits.b2.val[3], y.val[1]); + auto dot34 = vpaddq_s32(dot3, dot4); // 2, 2, 3, 3 + auto dot = vpaddq_s32(dot12, dot34); // 0, 1, 2, 3 + auto d8 = vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][2*i+0].d)); + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(scales.val[0], d8), vcvtq_f32_s32(dot)); + } + + deq.prepare(i, 1); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][2*i+1].qs); + auto dot1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[0], y.val[0]), deq.bits.b1.val[1], y.val[1]); + y = vld1q_s8_x2(q8.y[iy][2*i+1].qs+32); + auto dot2 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b1.val[2], y.val[0]), deq.bits.b1.val[3], y.val[1]); + auto dot12 = vpaddq_s32(dot1, dot2); // 0, 0, 1, 1 + y = vld1q_s8_x2(q8.y[iy][2*i+1].qs+64); + auto dot3 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[0], y.val[0]), deq.bits.b2.val[1], y.val[1]); + y = vld1q_s8_x2(q8.y[iy][2*i+1].qs+96); + auto dot4 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b2.val[2], y.val[0]), deq.bits.b2.val[3], y.val[1]); + auto dot34 = vpaddq_s32(dot3, dot4); // 2, 2, 3, 3 + auto dot = vpaddq_s32(dot12, dot34); // 0, 1, 2, 3 + auto d8 = vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][2*i+1].d)); + acc[iy] = vfmaq_f32(acc[iy], vmulq_f32(scales.val[1], d8), vcvtq_f32_s32(dot)); + } + } + + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, vaddvq_f32(acc[iy])); + } + } +} + +inline float convert_to_q8_k_r8(float d0, const int8x16x2_t * qx, const int8_t * scales, uint32_t * block, uint32_t * q8_k) { + auto max_i16 = vdupq_n_u16(0); + int16x8x4_t q[8]; + for (int ib32 = 0; ib32 < 8; ++ib32) { + auto scale_l = vdup_n_s8(scales[2*ib32+0]); + auto scale_h = vdup_n_s8(scales[2*ib32+1]); + q[ib32].val[0] = vmull_s8(scale_l, vget_low_s8 (qx[ib32].val[0])); + q[ib32].val[1] = vmull_s8(scale_l, vget_high_s8(qx[ib32].val[0])); + q[ib32].val[2] = vmull_s8(scale_h, vget_low_s8 (qx[ib32].val[1])); + q[ib32].val[3] = vmull_s8(scale_h, vget_high_s8(qx[ib32].val[1])); + max_i16 = vmaxq_u16(max_i16, vmaxq_u16(vabsq_s16(q[ib32].val[0]), vabsq_s16(q[ib32].val[1]))); + max_i16 = vmaxq_u16(max_i16, vmaxq_u16(vabsq_s16(q[ib32].val[2]), vabsq_s16(q[ib32].val[3]))); + } + uint16_t imax = vmaxvq_u16(max_i16); + if (!imax) { + for (int ib32 = 0; ib32 < 8; ++ib32) for (int l = 0; l < 8; ++l) q8_k[64*ib32 + 8*l] = 0; + return 0.f; + } + float dnew = float(imax) * d0; + //auto max_u32 = vmaxq_u32(vmovl_u16(vget_low_u16(max_i16)), vmovl_u16(vget_high_u16(max_i16))); + //auto max_f32 = vcvtq_f32_u32(max_u32); + //auto dnew = vmaxvq_f32(max_f32) * d0; + bool needs_scaling = true; + if (dnew <= 1.f) { + dnew = 1.f; needs_scaling = false; + } + auto scale = vdupq_n_f32(1/dnew); + for (int ib32 = 0; ib32 < 8; ++ib32) { + if (needs_scaling) { + for (int l = 0; l < 4; ++l) { + auto i1 = vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_low_s16 (q[ib32].val[l]))))); + auto i2 = vcvtnq_s32_f32(vmulq_f32(scale, vcvtq_f32_s32(vmovl_s16(vget_high_s16(q[ib32].val[l]))))); + q[ib32].val[l] = vcombine_s16(vmovn_s32(i1), vmovn_s32(i2)); + } + } + for (int l = 0; l < 2; ++l) { + auto s8 = vcombine_s8(vmovn_s16(q[ib32].val[2*l+0]), vmovn_s16(q[ib32].val[2*l+1])); + vst1q_s8((int8_t *)block + 16*l, s8); + } + auto qb = q8_k + 64*ib32; + for (int l = 0; l < 8; ++l) { + qb[8*l] = block[l]; + } + } + return dnew; +} + +// TODO: move this to iqk_gemm_iquants +void iqk_convert_iq4_xs_q8_k_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK_K; + + const block_iq4_xs * x8[8]; + + block_q8_k_r8 * y = (block_q8_k_r8 *)vy; + + auto values = vld1q_s8(iq4k_values); + + int8_t ls[16]; + float dnew[8]; + uint32_t block[8]; + int8x16x2_t xv[8]; + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const block_iq4_xs *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + float d = GGML_FP16_TO_FP32(x8[k][i].d); + for (int ib32 = 0; ib32 < 8; ++ib32) { + ls[2*ib32+0] = ls[2*ib32+1] = (((x8[k][i].scales_l[ib32/2] >> 4*(ib32%2)) & 0xf) | (((x8[k][i].scales_h >> 2*ib32) & 3) << 4)) - 32; + auto bits = vld1q_u8(x8[k][i].qs + 16*ib32); + xv[ib32].val[0] = vqtbl1q_s8(values, vandq_u8(bits, vdupq_n_u8(0xf))); + xv[ib32].val[1] = vqtbl1q_s8(values, vshrq_n_u8(bits, 4)); + } + dnew[k] = d * convert_to_q8_k_r8(1.f/127, xv, ls, block, (uint32_t *)y[i].qs + k); + } + vst1_f16((float16_t *)y[i].d + 0, vcvt_f16_f32(vld1q_f32(dnew+0))); + vst1_f16((float16_t *)y[i].d + 4, vcvt_f16_f32(vld1q_f32(dnew+4))); + } + y += nb; + } +} + +} + +bool iqk_convert_kquants_q8X_r8(int type, int n, const void * vx, size_t bx, void * vy, int nrc_x) { + switch (ggml_type(type)) { + case GGML_TYPE_Q2_K: iqk_convert_q2_k_q8_k_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q3_K: iqk_convert_q3_k_q8_k_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q4_K: iqk_convert_q4_k_q8_1_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q5_K: iqk_convert_q5_k_q8_1_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q6_K: iqk_convert_q6_k_q8_0_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_IQ4_XS: iqk_convert_iq4_xs_q8_k_r8(n, vx, bx, vy, nrc_x); break; + default: return false; + } + return true; +} + +bool iqk_set_kernels_kquants(int ne00, int typeA, int typeB, std::array& kernels, [[maybe_unused]] mul_mat_t& func16) { + + auto etypeA = ggml_type(typeA); + auto expected_type_B = etypeA == GGML_TYPE_IQ4_XS_R8 || etypeA == GGML_TYPE_Q4_K_R4 || etypeA == GGML_TYPE_Q5_K_R4 ? GGML_TYPE_Q8_K32 + //: etypeA == GGML_TYPE_Q8_K_R8 ? GGML_TYPE_Q8_KR8 + : etypeA == GGML_TYPE_Q8_KV || etypeA == GGML_TYPE_Q8_KV_R8 ? GGML_TYPE_Q8_KV + : etypeA == GGML_TYPE_Q6_K ? GGML_TYPE_Q8_0_X4 + : etypeA == GGML_TYPE_Q4_K || etypeA == GGML_TYPE_Q5_K ? GGML_TYPE_Q8_1_X4 + : GGML_TYPE_Q8_K; + + if (ne00%QK_K != 0 || ggml_type(typeB) != expected_type_B) { + return false; + } + + func16 = nullptr; + + switch (typeA) { + case GGML_TYPE_Q2_K: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerQ2K, kernels) + break; + case GGML_TYPE_Q3_K: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerQ3K, kernels) + break; + case GGML_TYPE_Q4_K: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_k_q8_1_x4, DequantizerQ4K, kernels) + break; + case GGML_TYPE_Q5_K: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_k_q8_1_x4, DequantizerQ5K, kernels) + break; + case GGML_TYPE_Q6_K: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q6_k_q8_0_x4, kernels) + break; + case GGML_TYPE_IQ4_XS: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_K_q8_K_T, DequantizerIQ4XS, kernels) + break; + case GGML_TYPE_Q2_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q2_k_r4_q8_k, kernels) + break; + case GGML_TYPE_Q3_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q3_k_r4_q8_k, kernels) + break; + case GGML_TYPE_Q4_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q4_k_r4_q8_k, kernels) + break; + case GGML_TYPE_Q5_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q5_k_r4_q8_k, kernels) + break; + case GGML_TYPE_Q6_K_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q6_k_r4_q8_k, kernels) + break; + case GGML_TYPE_IQ4_XS_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_iq4_xs_r8_q8_k, kernels) + break; + case GGML_TYPE_Q8_K_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_k_r8_q8_k, kernels) + break; + case GGML_TYPE_Q8_KV: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_KV_q8_KV, kernels) + kernels[0] = mul_mat_q8_KV_q8_KV_1; + func16 = mul_mat_q8_KV_q8_KV<16>; + break; + case GGML_TYPE_Q8_KV_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_KV_r8_q8_KV, kernels); + break; + default: + return false; + } + + return true; + +} + +#endif + +namespace { + +#ifdef __AVX2__ +template +static void mul_mat_q8_KV_q8_KV_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(n%32 == 0); + if (nrc_y == 1 && nrc_x == 1) { + auto dx = (const float *)vx; + auto dy = (const float *)info.src1_row(0); +#ifdef HAVE_FANCY_SIMD + auto sy = (const int32_t *)(dy + 1); + auto x = (const int8_t *)(dx + 2); + auto y = (const int8_t *)(dy + 2); + auto isum = _mm512_setzero_si512(); + for (int i = 0; i < n/64; ++i) { + auto qx = _mm512_loadu_si512((const __m512i *)x + i); + auto qy = _mm512_loadu_si512((const __m512i *)y + i); + isum = _mm512_dpbusd_epi32(isum, _mm512_add_epi8(qx, _mm512_set1_epi8(127)), qy); + } + auto isum256 = _mm256_add_epi32(_mm512_castsi512_si256(isum), _mm512_extracti32x8_epi32(isum, 1)); + for (int i = 2*(n/64); i < n/32; ++i) { + auto qx = _mm256_loadu_si256((const __m256i *)x + i); + auto qy = _mm256_loadu_si256((const __m256i *)y + i); + isum256 = _mm256_dpbusd_epi32(isum256, _mm256_add_epi8(qx, _mm256_set1_epi8(127)), qy); + } + info.store(0, 0, dx[0]*dy[0]*(hsum_i32_8(isum256) - 127*sy[0])); +#else + auto x = (const int8_t *)(dx + 2); + auto y = (const int8_t *)(dy + 2); + auto isum = _mm256_setzero_si256(); + for (int i = 0; i < n/32; ++i) { + auto qx = _mm256_loadu_si256((const __m256i *)x + i); + auto qy = _mm256_loadu_si256((const __m256i *)y + i); + auto dot = _mm256_maddubs_epi16(_mm256_sign_epi8(qx, qx), _mm256_sign_epi8(qy, qx)); + isum = _mm256_add_epi32(isum, _mm256_madd_epi16(_mm256_set1_epi16(1), dot)); + } + info.store(0, 0, dx[0]*dy[0]*hsum_i32_8(isum)); +#endif + return; + } + __m256i qx[2]; + __m256i acc[2*nrc_y] = {}; + float dy[nrc_y]; +#ifdef HAVE_FANCY_SIMD + int32_t sy[nrc_y]; +#else + __m256i sx[2]; + auto m1 = _mm256_set1_epi16(1); +#endif + const int8_t * q8y[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) { + auto dptr = (const float *)info.src1_row(iy); + dy[iy] = dptr[0]; +#ifdef HAVE_FANCY_SIMD + auto iptr = (const int32_t *)(dptr+1); + sy[iy] = -127*iptr[0]; +#endif + q8y[iy] = (const int8_t *)(dptr + 2); + } + for (int ix = 0; ix < nrc_x; ++ix) { + auto dx = (const float *)((const char *)vx + ix*bx); + auto q8x = (const int8_t *)(dx + 2); + for (int i = 0; i < n/64; ++i) { + for (int j = 0; j < 2; ++j) { +#ifdef HAVE_FANCY_SIMD + qx[j] = _mm256_add_epi8(_mm256_loadu_si256((const __m256i *)q8x + 2*i + j), _mm256_set1_epi8(127)); +#else + qx[j] = _mm256_loadu_si256((const __m256i *)q8x + 2*i + j); + sx[j] = _mm256_sign_epi8(qx[j], qx[j]); +#endif + } + for (int iy = 0; iy < nrc_y; ++iy) { + for (int j = 0; j < 2; ++j) { +#ifdef HAVE_FANCY_SIMD + acc[2*iy+j] = _mm256_dpbusd_epi32(acc[2*iy+j], qx[j], _mm256_loadu_si256((const __m256i *)q8y[iy] + 2*i + j)); +#else + auto dot = _mm256_maddubs_epi16(sx[j], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i *)q8y[iy] + 2*i + j), qx[j])); + acc[2*iy+j] = _mm256_add_epi32(acc[2*iy+j], _mm256_madd_epi16(m1, dot)); +#endif + } + } + } + if (int i = 2*(n/64); i < n/32) { +#ifdef HAVE_FANCY_SIMD + qx[0] = _mm256_add_epi8(_mm256_loadu_si256((const __m256i *)q8x + i), _mm256_set1_epi8(127)); +#else + qx[0] = _mm256_loadu_si256((const __m256i *)q8x + i); + sx[0] = _mm256_sign_epi8(qx[0], qx[0]); +#endif + for (int iy = 0; iy < nrc_y; ++iy) { +#ifdef HAVE_FANCY_SIMD + acc[2*iy] = _mm256_dpbusd_epi32(acc[2*iy], qx[0], _mm256_loadu_si256((const __m256i *)q8y[iy] + i)); +#else + auto dot = _mm256_maddubs_epi16(sx[0], _mm256_sign_epi8(_mm256_loadu_si256((const __m256i *)q8y[iy] + i), qx[0])); + acc[2*iy] = _mm256_add_epi32(acc[2*iy], _mm256_madd_epi16(m1, dot)); +#endif + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = hsum_i32_8(_mm256_add_epi32(acc[2*iy], acc[2*iy+1])); +#ifdef HAVE_FANCY_SIMD + info.store(ix, iy, dx[0]*dy[iy]*(sumi+sy[iy])); +#else + info.store(ix, iy, dx[0]*dy[iy]*sumi); +#endif + acc[2*iy] = acc[2*iy+1] = _mm256_setzero_si256(); + } + } +} + +#ifdef HAVE_FANCY_SIMD +template +static void mul_mat_q8_KV_q8_KV_8(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + GGML_ASSERT(n%32 == 0); + __m512i qx[4]; + __m512i acc[nrc_y <= 4 ? 2*nrc_y : nrc_y] = {}; + float dy[nrc_y]; + int32_t sy[nrc_y]; + const int8_t * q8y[nrc_y]; + for (int iy = 0; iy < nrc_y; ++iy) { + auto dptr = (const float *)info.src1_row(iy); + dy[iy] = dptr[0]; + auto iptr = (const int32_t *)(dptr + 1); + sy[iy] = -64*iptr[0]; + q8y[iy] = (const int8_t *)(dptr + 2); + } + const int8_t * q8x[8]; + float dx[8]; + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int kx = 0; kx < 8; ++kx) { + auto dptr = (const float *)((const char *)vx + (ix+kx)*bx); + dx[kx] = dptr[0]; + q8x[kx] = (const int8_t *)(dptr + 2); + } + for (int i = 0; i < n/32; ++i) { + for (int kx = 0; kx < 4; ++kx) { + qx[kx] = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)q8x[kx+0] + i)), + _mm256_loadu_si256((const __m256i *)q8x[kx+4] + i), 1); + } + auto t0 = _mm512_unpacklo_epi32(qx[0], qx[1]); + auto t1 = _mm512_unpacklo_epi32(qx[2], qx[3]); + auto t2 = _mm512_unpackhi_epi32(qx[0], qx[1]); + auto t3 = _mm512_unpackhi_epi32(qx[2], qx[3]); + qx[0] = _mm512_xor_si512(_mm512_unpacklo_epi64(t0, t1), _mm512_set1_epi8(-128)); + qx[1] = _mm512_xor_si512(_mm512_unpackhi_epi64(t0, t1), _mm512_set1_epi8(-128)); + qx[2] = _mm512_xor_si512(_mm512_unpacklo_epi64(t2, t3), _mm512_set1_epi8(-128)); + qx[3] = _mm512_xor_si512(_mm512_unpackhi_epi64(t2, t3), _mm512_set1_epi8(-128)); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y256 = _mm256_loadu_si256((const __m256i *)q8y[iy] + i); + auto y = _mm512_inserti32x8(_mm512_castsi256_si512(y256), y256, 1); + if constexpr (nrc_y <= 4) { + acc[2*iy+0] = _mm512_dpbusd_epi32(acc[2*iy+0], qx[0], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + acc[2*iy+1] = _mm512_dpbusd_epi32(acc[2*iy+1], qx[1], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + acc[2*iy+0] = _mm512_dpbusd_epi32(acc[2*iy+0], qx[2], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + acc[2*iy+1] = _mm512_dpbusd_epi32(acc[2*iy+1], qx[3], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + } else { + acc[iy] = _mm512_dpbusd_epi32(acc[iy], qx[0], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + acc[iy] = _mm512_dpbusd_epi32(acc[iy], qx[1], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + acc[iy] = _mm512_dpbusd_epi32(acc[iy], qx[2], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + acc[iy] = _mm512_dpbusd_epi32(acc[iy], qx[3], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + } + } + } + auto scales_x = _mm256_loadu_ps(dx); + for (int iy = 0; iy < nrc_y; ++iy) { + if constexpr (nrc_y <= 4) { + auto ss = _mm512_add_epi32(_mm512_add_epi32(acc[2*iy+0], acc[2*iy+1]), _mm512_set1_epi32(sy[iy])); + auto sum1 = _mm_add_epi32(_mm512_extracti32x4_epi32(ss, 0), _mm512_extracti32x4_epi32(ss, 1)); + auto sum2 = _mm_add_epi32(_mm512_extracti32x4_epi32(ss, 2), _mm512_extracti32x4_epi32(ss, 3)); + auto scale = _mm256_mul_ps(scales_x, _mm256_set1_ps(dy[iy])); + info.store(ix+0, iy, _mm_mul_ps(_mm256_castps256_ps128(scale), _mm_cvtepi32_ps(sum1))); + info.store(ix+4, iy, _mm_mul_ps(_mm256_extractf128_ps(scale, 1), _mm_cvtepi32_ps(sum2))); + acc[2*iy+0] = acc[2*iy+1] = _mm512_setzero_si512(); + } else { + acc[iy] = _mm512_add_epi32(acc[iy], _mm512_set1_epi32(sy[iy])); + auto sum1 = _mm_add_epi32(_mm512_extracti32x4_epi32(acc[iy], 0), _mm512_extracti32x4_epi32(acc[iy], 1)); + auto sum2 = _mm_add_epi32(_mm512_extracti32x4_epi32(acc[iy], 2), _mm512_extracti32x4_epi32(acc[iy], 3)); + auto scale = _mm256_mul_ps(scales_x, _mm256_set1_ps(dy[iy])); + info.store(ix+0, iy, _mm_mul_ps(_mm256_castps256_ps128(scale), _mm_cvtepi32_ps(sum1))); + info.store(ix+4, iy, _mm_mul_ps(_mm256_extractf128_ps(scale, 1), _mm_cvtepi32_ps(sum2))); + acc[iy] = _mm512_setzero_si512(); + } + } + } +} +#endif +#endif + +template +inline std::pair mul_mat_kernel([[maybe_unused]] int D, int int_typeA, int nq) { + auto typeA = ggml_type(int_typeA); + constexpr int kMaxQ = 8; +#define MAKE_FUNCS(mul_mat, n) \ + if (n >= kMaxQ) return std::make_pair(mul_mat, kMaxQ>, kMaxQ);\ + else {\ + switch (n) {\ + case 1: return std::make_pair(mul_mat, 1>, 1);\ + case 2: return std::make_pair(mul_mat, 2>, 2);\ + case 3: return std::make_pair(mul_mat, 3>, 3);\ + case 4: return std::make_pair(mul_mat, 4>, 4);\ + case 5: return std::make_pair(mul_mat, 5>, 5);\ + case 6: return std::make_pair(mul_mat, 6>, 6);\ + case 7: return std::make_pair(mul_mat, 7>, 7);\ + }\ + } +#define MAKE_FUNCS_ONLY_NRC(mul_mat, n) \ + if (n >= kMaxQ) return std::make_pair(mul_mat, kMaxQ);\ + else {\ + switch (n) {\ + case 1: return std::make_pair(mul_mat<1>, 1);\ + case 2: return std::make_pair(mul_mat<2>, 2);\ + case 3: return std::make_pair(mul_mat<3>, 3);\ + case 4: return std::make_pair(mul_mat<4>, 4);\ + case 5: return std::make_pair(mul_mat<5>, 5);\ + case 6: return std::make_pair(mul_mat<6>, 6);\ + case 7: return std::make_pair(mul_mat<7>, 7);\ + }\ + } + if (typeA == GGML_TYPE_Q8_KV) { +#ifdef __aarch64__ + if (nq%16 == 0) return std::make_pair(mul_mat_q8_KV_q8_KV<16>, 16); + if (nq == 1) return std::make_pair(mul_mat_q8_KV_q8_KV_1, 1); + MAKE_FUNCS_ONLY_NRC(mul_mat_q8_KV_q8_KV, nq); +#else + if (nq == 1) return std::make_pair(mul_mat_q8_KV_q8_KV_1<1>, 1); +#ifdef HAVE_FANCY_SIMD + if (D%32 == 0 && k_step%8 == 0) { + if (nq%16 == 0) return std::make_pair(mul_mat_q8_KV_q8_KV_8<16>, 16); + MAKE_FUNCS_ONLY_NRC(mul_mat_q8_KV_q8_KV_8, nq); + } else { + if (nq%16 == 0) return std::make_pair(mul_mat_q8_KV_q8_KV<16>, 16); + } +#endif + MAKE_FUNCS_ONLY_NRC(mul_mat_q8_KV_q8_KV, nq); +#endif + } + else if (typeA == GGML_TYPE_Q8_KV_R8) { + MAKE_FUNCS_ONLY_NRC(mul_mat_q8_KV_r8_q8_KV, nq); + } + GGML_ABORT("Fatal error"); +} + +inline std::pair mul_mat_kernel(int D, int int_typeA, int nq, int k_step) { + switch (k_step) { + case 32: return mul_mat_kernel< 32>(D, int_typeA, nq); + case 64: return mul_mat_kernel< 64>(D, int_typeA, nq); + case 128: return mul_mat_kernel<128>(D, int_typeA, nq); + default: GGML_ABORT("Fatal error"); + } +} + +} + +void iqk_gemm_q8kv_fa(int D, int nq, int type_k, const char * k, size_t stride_k, DataInfo& info, int k_step) { + auto [mul_mat, nrc_q] = mul_mat_kernel(D, type_k, nq, k_step); + for (int iq = 0; iq < nq/nrc_q; ++iq) { + mul_mat(D, k, stride_k, info, k_step); + info.cur_y += nrc_q; + } + int iq = nrc_q*(nq/nrc_q); + if (iq < nq) { + auto [mul_mat1, nrc_q1] = mul_mat_kernel(D, type_k, nq - iq, k_step); + GGML_ASSERT(nrc_q1 == nq - iq); + mul_mat1(D, k, stride_k, info, k_step); + } +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_gemm_kquants.h b/llama.cpp/ggml/src/iqk/iqk_gemm_kquants.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_gemm_kquants.h @@ -0,0 +1,15 @@ +#pragma once + +#include "iqk_common.h" + +#ifdef IQK_IMPLEMENT + +#include + +bool iqk_set_kernels_kquants(int ne00, int typeA, int typeB, std::array& kernels, mul_mat_t& func16); + +void iqk_gemm_q8kv_fa(int D, int nq, int type_k, const char * k, size_t stride_k, DataInfo& info, int k_step); + +bool iqk_convert_kquants_q8X_r8(int type, int n, const void * vx, size_t bx, void * vy, int nrc_x); + +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_gemm_ktquants.h b/llama.cpp/ggml/src/iqk/iqk_gemm_ktquants.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_gemm_ktquants.h @@ -0,0 +1,13 @@ +#pragma once + +#include "iqk_common.h" + +#ifdef IQK_IMPLEMENT + +#include + +bool iqk_set_kernels_ktquants(int ne00, int typeA, int typeB, std::array& kernels, mul_mat_t& func16); + +bool iqk_dequantize_ktquants(int type, int n, const void * vx, size_t bx, void * vy, size_t stride_y, int nrc_x); + +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_gemm_legacy_quants.cpp b/llama.cpp/ggml/src/iqk/iqk_gemm_legacy_quants.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_gemm_legacy_quants.cpp @@ -0,0 +1,3521 @@ +#include "iqk_gemm_legacy_quants.h" + +#include + +#ifdef IQK_IMPLEMENT + +#include "ggml-impl.h" +#include "iqk_utils.h" + +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" + +// +// ============================== Legacy quants +// + +#ifdef __x86_64__ + +namespace { + +struct DotHelper { + const __m256i m1 = _mm256_set1_epi16(1); +#if defined(__AVX512VNNI__) && defined(__AVX512VL__) + inline __m256i dot(__m256i x, __m256i y) const { + return _mm256_dpbusd_epi32(_mm256_setzero_si256(), x, y); + } +#else + inline __m256i dot(__m256i x, __m256i y) const { + return _mm256_madd_epi16(m1, _mm256_maddubs_epi16(x, y)); + } +#endif +}; + +struct SignedDot { + DotHelper helper; + inline __m256i compute(__m256i x, __m256i y) const { + return helper.dot(_mm256_sign_epi8(x, x), _mm256_sign_epi8(y, x)); + } +}; +struct UnsignedDot { + DotHelper helper; + inline __m256i compute(__m256i x, __m256i y) const { + return helper.dot(x, y); + } +}; + +template struct Sum4 { + Dot dot; + inline __m256i compute(const __m256i * qx, const Q8 * y) const { + const Q8x4 * y4 = (const Q8x4 *)y; + const __m256i p0 = dot.compute(qx[0], _mm256_loadu_si256((const __m256i *)y4->qs+0)); // 8x block 0 + const __m256i p1 = dot.compute(qx[1], _mm256_loadu_si256((const __m256i *)y4->qs+1)); // 8x block 1 + const __m256i p2 = dot.compute(qx[2], _mm256_loadu_si256((const __m256i *)y4->qs+2)); // 8x block 2 + const __m256i p3 = dot.compute(qx[3], _mm256_loadu_si256((const __m256i *)y4->qs+3)); // 8x block 3 + if constexpr (can_pack) { + const __m256i p01 = _mm256_madd_epi16(dot.helper.m1, _mm256_packs_epi32(p0, p1)); // 0,0, 1,1, 0,0, 1,1 + const __m256i p23 = _mm256_madd_epi16(dot.helper.m1, _mm256_packs_epi32(p2, p3)); // 2,2, 3,3, 2,2, 3,3 + return _mm256_madd_epi16(dot.helper.m1, _mm256_packs_epi32(p01, p23)); // 0,1,2,3, 0,1,2,3 + } else { + // Note to myself: this is much faster than using _mm256_hadd_epi32() + auto p01 = _mm256_add_epi32(_mm256_unpacklo_epi32(p0, p1), _mm256_unpackhi_epi32(p0, p1)); // 0,1, 0,1, 0,1, 0,1 + auto p23 = _mm256_add_epi32(_mm256_unpacklo_epi32(p2, p3), _mm256_unpackhi_epi32(p2, p3)); // 2,3, 2,3, 2,3, 2,3 + return _mm256_add_epi32(_mm256_unpacklo_epi64(p01, p23), _mm256_unpackhi_epi64(p01, p23)); // 0,1,2,3, 0,1,2,3 + } + } + inline __m256i compute(__m256i x, __m256i y) const { return dot.compute(x, y); } +}; + +template struct Sum4q4 { + inline __m256i compute(const __m256i * qx, const Q8 * y) const { + const Q8x4 * y4 = (const Q8x4 *)y; + auto p0 = _mm256_maddubs_epi16(qx[0], _mm256_loadu_si256((const __m256i *)y4->qs+0)); // 16x block 0 + auto p1 = _mm256_maddubs_epi16(qx[1], _mm256_loadu_si256((const __m256i *)y4->qs+1)); // 16x block 1 + auto p2 = _mm256_maddubs_epi16(qx[2], _mm256_loadu_si256((const __m256i *)y4->qs+2)); // 16x block 2 + auto p3 = _mm256_maddubs_epi16(qx[3], _mm256_loadu_si256((const __m256i *)y4->qs+3)); // 16x block 3 + auto p01 = _mm256_add_epi16(_mm256_unpacklo_epi32(p0, p1), _mm256_unpackhi_epi32(p0, p1)); // 0,0, 1,1, 0,0, 1,1, 0,0, 1,1, 0,0, 1,1 + auto p23 = _mm256_add_epi16(_mm256_unpacklo_epi32(p2, p3), _mm256_unpackhi_epi32(p2, p3)); // 2,2, 3,3, 2,2, 3,3, 2,2, 3,3, 2,2, 3,3 + auto p0123 = _mm256_add_epi16(_mm256_unpacklo_epi64(p01, p23), _mm256_unpackhi_epi64(p01, p23)); // 0,0, 1,1, 2,2, 3,3, 0,0, 1,1, 2,2, 3,3 + return _mm256_madd_epi16(_mm256_set1_epi16(1), p0123); + } + inline __m256i compute(__m256i x, __m256i y) const { return _mm256_madd_epi16(_mm256_set1_epi16(1), _mm256_maddubs_epi16(x, y)); } +}; + +inline __m256 convert_scales(const uint16_t * scales) { + auto aux_d = _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)scales)), 16)); + auto aux_m = _mm_cvtepi32_ps(_mm_cvtepi16_epi32(_mm_loadl_epi64((const __m128i *)(scales+4)))); + return _mm256_set_m128(_mm_mul_ps(aux_d, aux_m), aux_d); +} + +inline __m128 convert_scales_s(const uint16_t * scales) { + return _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)scales)), 16)); +} + +struct ScaleHelperQ8_0 { + inline __m128 prepare4(const block_q8_0 * y) { + const block_q8_0_x4 * y4 = (const block_q8_0_x4 *)y; + return _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)y4->d)); + } + inline __m128 prepare4(__m128 other_scales, const block_q8_0 * y) { + return _mm_mul_ps(other_scales, prepare4(y)); + } + template inline float prepare1(const Q * y) const { return GGML_FP16_TO_FP32(y->d); } + template inline float prepare1(float d, const Q * y) const { return d*prepare1(y); } +}; + +struct ScaleHelperQ_0 { + ggml_half scales8[4]; + template + inline __m128 prepare4(const Q * y) { + for (int j = 0; j < 4; ++j) scales8[j] = y[j].d; + return _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)scales8)); + } + template + inline __m128 prepare4(__m128 other_scales, const Q * y) { + return _mm_mul_ps(other_scales, prepare4(y)); + } + template inline float prepare1(const Q * y) const { return GGML_FP16_TO_FP32(y->d); } + template inline float prepare1(float d, const Q * y) const { return d*prepare1(y); } +}; + +struct ScaleHelperQ8_2S { + template + inline __m128 prepare4(const Q * y) { + const block_q8_2_x4 * y4 = (const block_q8_2_x4 *)y; + return convert_scales_s((const uint16_t *)y4->d); + } + template + inline __m128 prepare4(__m128 other_scales, const Q * y) { + return _mm_mul_ps(other_scales, prepare4(y)); + } + template static inline float prepare1(const Q * y) { return GGML_BF16_TO_FP32(ggml_bf16_t{y->d}); } + template static inline float prepare1(float d, const Q * y) { return d*prepare1(y); } +}; + +struct ScaleHelperQ_0_MXFP4 { + float scales[4]; + template + inline __m128 prepare4(const Q * y) { + for (int j = 0; j < 4; ++j) scales[j] = GGML_E8M0_TO_FP32_HALF(y[j].e); + return _mm_loadu_ps(scales); + } + template + inline __m128 prepare4(__m128 other_scales, const Q * y) { + return _mm_mul_ps(other_scales, prepare4(y)); + } + template inline float prepare1(const Q * y) const { return GGML_E8M0_TO_FP32_HALF(y->e); } + template inline float prepare1(float d, const Q * y) const { return d*prepare1(y); } +}; + +template +struct ScaleHelperQ_0_1 { + ggml_half scales8[4]; + template + inline __m256 prepare4(const Q * y) { + for (int j = 0; j < 4; ++j) scales8[j] = y[j].d; + auto s4 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)scales8)); + return _mm256_set_m128(_mm_mul_ps(s4, min), s4); + } + template + inline __m256 prepare4(__m256 other_scales, const Q * y) { + return _mm_mul256_ps(other_scales, prepare4(y)); + } + template inline std::pair prepare1(const Q * y) const { + float d = GGML_FP16_TO_FP32(y->d); + return std::make_pair(d, -d*float(min_value)); + } + std::pair inline prepare1(const std::pair& dm, const block_q8_1 * y) const { + return std::make_pair(dm.first*GGML_FP16_TO_FP32(y->d), dm.second*GGML_FP16_TO_FP32(y->s)); + } + const __m128 min = _mm_set1_ps(float(-min_value)); +}; + +template +struct ScaleHelperQ_0_1_MXFP4 { + float scales[4]; + template + inline __m256 prepare4(const Q * y) { + for (int j = 0; j < 4; ++j) scales[j] = GGML_E8M0_TO_FP32_HALF(y[j].e); + auto s4 = _mm_loadu_ps(scales); + return _mm256_set_m128(_mm_mul_ps(s4, min), s4); + } + template + inline __m256 prepare4(__m256 other_scales, const Q * y) { + return _mm_mul256_ps(other_scales, prepare4(y)); + } + template inline std::pair prepare1(const Q * y) const { + float d = GGML_E8M0_TO_FP32_HALF(y->e); + return std::make_pair(d, -d*float(min_value)); + } + std::pair inline prepare1(const std::pair& dm, const block_q8_1 * y) const { + return std::make_pair(dm.first*GGML_FP16_TO_FP32(y->d), dm.second*GGML_FP16_TO_FP32(y->s)); + } + const __m128 min = _mm_set1_ps(float(-min_value)); +}; + +struct ScaleHelperQ8_1 { + template + inline __m256 prepare4(const Q * y) { + const block_q8_1_x4 * y4 = (const block_q8_1_x4 *)y; + return _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)y4->d)); + } + template + inline __m256 prepare4(__m256 other_scales, const Q * y) { + return _mm256_mul_ps(other_scales, prepare4(y)); + } + template inline std::pair prepare1(const Q * y) const { + return std::make_pair(GGML_FP16_TO_FP32(y->d), GGML_FP16_TO_FP32(y->m)); + } + template inline std::pair prepare1(const std::pair& dm, const Q * y) const { + return std::make_pair(dm.first*GGML_FP16_TO_FP32(y->d), dm.second*GGML_FP16_TO_FP32(y->m)); + } + std::pair inline prepare1(const std::pair& dm, const block_q8_1 * y) const { + return std::make_pair(dm.first*GGML_FP16_TO_FP32(y->d), dm.second*GGML_FP16_TO_FP32(y->s)); + } +}; + +struct ScaleHelperQ8_2 { + template + inline __m256 prepare4(const Q * y) { + const block_q8_2_x4 * y4 = (const block_q8_2_x4 *)y; + return convert_scales((const uint16_t *)y4->d); + } + template + inline __m256 prepare4(__m256 other_scales, const Q * y) { + return _mm256_mul_ps(other_scales, prepare4(y)); + } + template static inline std::pair prepare1(const Q * y) { + float d = GGML_BF16_TO_FP32(ggml_bf16_t{y->d}); + int16_t m = *(const int16_t *)&y->s; + return std::make_pair(d, d*m); + } + static inline std::pair prepare1(const std::pair& dm, const block_q8_2 * y) { + auto d = prepare1(y); + return std::make_pair(dm.first*d.first, dm.second*d.second); + } +}; + +struct ScaleHelperQ_1 { + uint32_t scales8[4]; + const __m128i shuffle = _mm_set_epi16(0x0f0e, 0x0b0a, 0x0706, 0x0302, 0x0d0c, 0x0908, 0x0504, 0x0100); + + template + inline __m256 prepare4(const Q * y) { + for (int j = 0; j < 4; ++j) { + // it is slightly faster to directly dereference (const uint32 *)&y[j].d, but some compilers + // complain that this breaks strict-aliasing rules. + memcpy(scales8 + j, &y[j].d, sizeof(uint32_t)); + } + return _mm256_cvtph_ps(_mm_shuffle_epi8(_mm_loadu_si128((const __m128i *)scales8), shuffle)); + } + + template + inline __m256 prepare4(__m256 other_scales, const Q * y) { + return _mm256_mul_ps(other_scales, prepare4(y)); + } + + template inline std::pair prepare1(const Q * y) const { + return std::make_pair(GGML_FP16_TO_FP32(y->d), GGML_FP16_TO_FP32(y->m)); + } + template inline std::pair prepare1(const std::pair& dm, const Q * y) const { + return std::make_pair(dm.first*GGML_FP16_TO_FP32(y->d), dm.second*GGML_FP16_TO_FP32(y->m)); + } + std::pair inline prepare1(const std::pair& dm, const block_q8_1 * y) const { + return std::make_pair(dm.first*GGML_FP16_TO_FP32(y->d), dm.second*GGML_FP16_TO_FP32(y->s)); + } +}; + +struct MinusType0 { + inline __m256 compute(__m128 d, int) const { return _mm256_set_m128(d, d); } + inline float compute(float d, int) const { return d; } + inline float result(__m256 acc, int) const { return hsum_float_8(acc); } + inline __m256 vresult(__m256 acc, int) const { return acc; } +}; + +template struct MinusType1 { + __m128 accm[nrc_y]; + MinusType1() { for (int iy = 0; iy < nrc_y; ++iy) accm[iy] = _mm_setzero_ps(); } + inline __m256 compute(__m256 dm, int iy) { + const __m128 d = _mm256_castps256_ps128(dm); + const __m128 m = _mm256_extractf128_ps(dm, 1); + accm[iy] = _mm_add_ps(accm[iy], m); + return _mm256_set_m128(d, d); + } + inline float compute(const std::pair& dm, int iy) { + accm[iy] = _mm_add_ps(accm[iy], _mm_set1_ps(dm.second*0.25f)); + return dm.first; + } + inline float result(__m256 acc, int iy) const { + const __m128 sum = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1)); + return hsum_float_4(_mm_add_ps(sum, accm[iy])); + } + inline __m256 vresult(__m256 acc, int iy) const { + return _mm256_add_ps(acc, _mm256_insertf128_ps(_mm256_setzero_ps(), accm[iy], 0)); + } +}; + +template struct AccumT { + __m256 acc[nrc_y]; + Minus accm; + AccumT() { for (int iy = 0; iy < nrc_y; ++iy) acc[iy] = _mm256_setzero_ps(); } + template + inline void compute(int nb, Unpacker& unp, Scales& scales, Sum& sum, const Q8 ** y, const DataInfo& info, int ix) { + auto qx = unp.quants(); + __m256 dall[nrc_y]; + for (int i = 0; i < nb/4; ++i) { + auto other_scales = unp.set_block_4(i); + for (int iy = 0; iy < nrc_y; ++iy) { + auto s12 = scales.prepare4(other_scales, y[iy] + 4*i); + dall[iy] = accm.compute(s12, iy); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto pall = sum.compute(qx, y[iy] + 4*i); + acc[iy] = _mm256_fmadd_ps(dall[iy], _mm256_cvtepi32_ps(pall), acc[iy]); + } + } + if (!is_multiple_of_4) { + for (int i = 4*(nb/4); i < nb; ++i) { + auto other_scales = unp.set_block(i); + for (int iy = 0; iy < nrc_y; ++iy) { + auto s12 = scales.prepare1(other_scales, y[iy] + i); + auto d = accm.compute(s12, iy); + const __m256i p0 = sum.compute(qx[0], _mm256_loadu_si256((const __m256i *)y[iy][i].qs)); + acc[iy] = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(p0), acc[iy]); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, accm.result(acc[iy], iy)); + } + } + template + inline void compute(int nb, Unpacker& unp, Scales& scales, Sum& sum, const Q8 ** y, __m256 * result) { + auto qx = unp.quants(); + __m256 dall[nrc_y]; + for (int i = 0; i < nb/4; ++i) { + auto other_scales = unp.set_block_4(i); + for (int iy = 0; iy < nrc_y; ++iy) { + auto s12 = scales.prepare4(other_scales, y[iy] + 4*i); + dall[iy] = accm.compute(s12, iy); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto pall = sum.compute(qx, y[iy] + 4*i); + acc[iy] = _mm256_fmadd_ps(dall[iy], _mm256_cvtepi32_ps(pall), acc[iy]); + } + } + if (!is_multiple_of_4) { + for (int i = 4*(nb/4); i < nb; ++i) { + auto other_scales = unp.set_block(i); + for (int iy = 0; iy < nrc_y; ++iy) { + auto s12 = scales.prepare1(other_scales, y[iy] + i); + auto d = accm.compute(s12, iy); + const __m256i p0 = sum.compute(qx[0], _mm256_loadu_si256((const __m256i *)y[iy][i].qs)); + acc[iy] = _mm256_fmadd_ps(_mm256_set1_ps(d), _mm256_cvtepi32_ps(p0), acc[iy]); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + result[iy] = accm.vresult(acc[iy], iy); + } + } +}; + +template +using AccumType0 = AccumT; + +template +using AccumType1 = AccumT, nrc_y, is_multiple_of_4>; + +using Sum4TypeQ80 = Sum4; +using Sum4TypeQ82 = Sum4; +using Sum4TypeQ82S = Sum4; + +template +void mul_mat_qX_q8_Helper(int nb, const void * vx, size_t bx, const DataInfo& info, const Q8 ** y, int nrc_x) { + Unpacker unp(vx, bx); + typename Unpacker::Sum4T sum4; + Scales scales; + for (int ix = 0; ix < nrc_x; ++ix) { + unp.set_row(ix); + AccumType accum; + accum.compute(nb, unp, scales, sum4, y, info, ix); + } +} + +template +void mul_mat_qX_q8_Helper_x2(int nb, const void * vx, size_t bx, const DataInfo& info, const Q8 ** y, int nrc_x) { + GGML_ASSERT(nrc_x%2 == 0); + Unpacker unp(vx, bx); + typename Unpacker::Sum4T sum4; + Scales scales; + for (int ix = 0; ix < nrc_x; ix += 2) { + unp.set_row(ix); + AccumType accum; + accum.compute(nb, unp, scales, sum4, y, info, ix); + } +} + +template +void mul_mat_qX_0_q8_0_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n%Unpacker::block_size() == 0); + Q8 q8(info); + int nb = n/Unpacker::block_size(); + if constexpr (std::is_same_v) { + if (nb%4 == 0) { + mul_mat_qX_q8_Helper, ScaleHelperQ8_2S, Block, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x); + } else { + mul_mat_qX_q8_Helper, ScaleHelperQ8_2S, Block, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x); + } + } + else { + if (nb%4 == 0) { + mul_mat_qX_q8_Helper, ScaleHelperQ8_0, Block, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x); + } else { + mul_mat_qX_q8_Helper, ScaleHelperQ8_0, Block, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x); + } + } +} + +template +void mul_mat_qX_0_q8_2_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n%Unpacker::block_size() == 0); + Q8 q8(info); + int nb = n/Unpacker::block_size(); + if (nb%4 == 0) { + mul_mat_qX_q8_Helper, ScaleHelperQ8_0, block_q8_0, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x + ); + } else { + mul_mat_qX_q8_Helper, ScaleHelperQ8_0, block_q8_0, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x + ); + } +} + +template +void mul_mat_qX_0_q8_0_Tx(int n, const void * vx, size_t bx, const DataInfo& info, int) { + static_assert(8%nrc_y == 0); + Q8 q8(info); + int nb = n/Unpacker::block_size(); + Unpacker unp(vx, bx); + typename Unpacker::Sum4T sum4; + ScaleHelperQ8_2S scales; + __m256 result[8]; + auto store = [&info, &result] (int ix0) { + if constexpr (nrc_y == 1) { + info.store(ix0, 0, hsum_float_8x8(result)); + } + else if constexpr (nrc_y == 2) { + auto value = hsum_float_8x8(result); + auto value1 = _mm256_extractf128_ps(value, 1); + info.store(ix0, 0, _mm_shuffle_ps(_mm256_castps256_ps128(value), value1, 0x88)); + info.store(ix0, 1, _mm_shuffle_ps(_mm256_castps256_ps128(value), value1, 0xdd)); + } + else { + float val[8]; + _mm256_storeu_ps(val, hsum_float_8x8(result)); + for (int iy = 0; iy < nrc_y; ++iy) for (int ix = 0; ix < 8/nrc_y; ++ix) info.store(ix0+ix, iy, val[nrc_y*ix+iy]); + } + }; + if (nb%4 == 0) { + for (int ix0 = 0; ix0 < nrc_x; ix0 += 8/nrc_y) { + for (int ix = 0; ix < 8/nrc_y; ++ix) { + unp.set_row(ix0 + ix); + AccumType0 accum; + accum.compute(nb, unp, scales, sum4, q8.y, result + nrc_y*ix); + } + store(ix0); + } + } else { + for (int ix0 = 0; ix0 < nrc_x; ix0 += 8/nrc_y) { + for (int ix = 0; ix < 8/nrc_y; ++ix) { + unp.set_row(ix0 + ix); + AccumType0 accum; + accum.compute(nb, unp, scales, sum4, q8.y, result + nrc_y*ix); + } + store(ix0); + } + } +} + +template +void mul_mat_qX_1_q8_1_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n%Unpacker::block_size() == 0); + Q8 q8(info); + int nb = n/Unpacker::block_size(); + if (nb%4 == 0) { + mul_mat_qX_q8_Helper, ScaleHelperQ8_1, block_q8_1, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x + ); + } else { + mul_mat_qX_q8_Helper, ScaleHelperQ8_1, block_q8_1, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x + ); + } +} + +template +void mul_mat_qX_1_q8_2_T(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + assert(n%Unpacker::block_size() == 0); + Q8 q8(info); + int nb = n/Unpacker::block_size(); + if (nb%4 == 0) { + mul_mat_qX_q8_Helper, ScaleHelperQ8_2, block_q8_2, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x + ); + } else { + mul_mat_qX_q8_Helper, ScaleHelperQ8_2, block_q8_2, nrc_y>( + nb, vx, bx, info, q8.y, nrc_x + ); + } +} + +template +void mul_mat_qX_0_q8_2_Tx(int n, const void * vx, size_t bx, const DataInfo& info, int) { + static_assert(8%nrc_y == 0); + Q8 q8(info); + int nb = n/Unpacker::block_size(); + Unpacker unp(vx, bx); + typename Unpacker::Sum4T sum4; + ScaleHelperQ8_2 scales; + __m256 result[8]; + auto store = [&info, &result] (int ix0) { + if constexpr (nrc_y == 1) { + info.store(ix0, 0, hsum_float_8x8(result)); + } + else if constexpr (nrc_y == 2) { + auto value = hsum_float_8x8(result); + auto value1 = _mm256_extractf128_ps(value, 1); + info.store(ix0, 0, _mm_shuffle_ps(_mm256_castps256_ps128(value), value1, 0x88)); + info.store(ix0, 1, _mm_shuffle_ps(_mm256_castps256_ps128(value), value1, 0xdd)); + } + else { + float val[8]; + _mm256_storeu_ps(val, hsum_float_8x8(result)); + for (int iy = 0; iy < nrc_y; ++iy) for (int ix = 0; ix < 8/nrc_y; ++ix) info.store(ix0+ix, iy, val[nrc_y*ix+iy]); + } + }; + if (nb%4 == 0) { + for (int ix0 = 0; ix0 < nrc_x; ix0 += 8/nrc_y) { + for (int ix = 0; ix < 8/nrc_y; ++ix) { + unp.set_row(ix0 + ix); + AccumType1 accum; + accum.compute(nb, unp, scales, sum4, q8.y, result + nrc_y*ix); + } + store(ix0); + } + } else { + for (int ix0 = 0; ix0 < nrc_x; ix0 += 8/nrc_y) { + for (int ix = 0; ix < 8/nrc_y; ++ix) { + unp.set_row(ix0 + ix); + AccumType1 accum; + accum.compute(nb, unp, scales, sum4, q8.y, result + nrc_y*ix); + } + store(ix0); + } + } +} + +struct Dequantizer4bit { + const __m256i m4 = _mm256_set1_epi8(0xf); + inline __m256i dequant(const uint8_t * qs) const { + const __m128i aux128 = _mm_loadu_si128((const __m128i *)qs); + return _mm256_and_si256(MM256_SET_M128I(_mm_srli_epi16(aux128, 4), aux128), m4); + } +}; + +struct Q8_0_Dequantizer { + inline __m256i dequant(const block_q8_0 * x) const { + return _mm256_loadu_si256((const __m256i *)x->qs); + } +}; + +struct Q8_0_1_Dequantizer { + inline __m256i dequant(const block_q8_0 * x) const { + return _mm256_add_epi8(_mm256_set1_epi8(127), _mm256_loadu_si256((const __m256i *)x->qs)); + } +}; + +struct Q4_0_Dequantizer { + Dequantizer4bit b4; + const __m256i m8 = _mm256_set1_epi8(-8); + inline __m256i dequant(const block_q4_0 * x) const { + return _mm256_add_epi8(b4.dequant(x->qs), m8); + } +}; + +struct Q4_0_1_Dequantizer { + Dequantizer4bit b4; + inline __m256i dequant(const block_q4_0 * x) const { + return b4.dequant(x->qs); + } +}; + +struct IQ4_NL_DequantizerU { + Dequantizer4bit b4; + const __m256i values = load_iq4nl_values_256(); + inline __m256i dequant(const block_iq4_nl * x) const { + return _mm256_shuffle_epi8(values, b4.dequant(x->qs)); + } +}; + +struct IQ4_NL_DequantizerS { + Dequantizer4bit b4; + const __m256i values = load_iq4k_values_256(); + inline __m256i dequant(const block_iq4_nl * x) const { + return _mm256_shuffle_epi8(values, b4.dequant(x->qs)); + } +}; + +//============================= +static inline __m128i load_unsigned_mxfp4_values_128() { + static const uint8_t kvalues_mxfp4_unsigned[16] = {12, 13, 14, 15, 16, 18, 20, 24, 12, 11, 10, 9, 8, 6, 4, 0}; + return _mm_loadu_si128((const __m128i *)kvalues_mxfp4_unsigned); +} + +static inline __m256i load_unsigned_mxfp4_values_256() { + auto val128 = load_unsigned_mxfp4_values_128(); + return MM256_SET_M128I(val128, val128); +} + +#ifdef HAVE_FANCY_SIMD +static inline __m512i load_unsigned_mxfp4_values_512() { + auto val256 = load_unsigned_mxfp4_values_256(); + return _mm512_inserti32x8(_mm512_castsi256_si512(val256), val256, 1); +} +#endif + +static inline __m128i load_mxfp4_values_128() { + return _mm_loadu_si128((const __m128i *)kvalues_mxfp4); +} + +static inline __m256i load_mxfp4_values_256() { + auto val128 = load_mxfp4_values_128(); + return MM256_SET_M128I(val128, val128); +} + +struct MXFP4_Dequantizer { + Dequantizer4bit b4; + const __m256i values = load_unsigned_mxfp4_values_256(); + inline __m256i dequant(const block_mxfp4 * x) const { + return _mm256_shuffle_epi8(values, b4.dequant(x->qs)); + } +}; + +struct MXFP40_Dequantizer { + Dequantizer4bit b4; + const __m256i values = load_mxfp4_values_256(); + inline __m256i dequant(const block_mxfp4 * x) const { + return _mm256_shuffle_epi8(values, b4.dequant(x->qs)); + } +}; + +struct Q4_1_Dequantizer { + Dequantizer4bit b4; + inline __m256i dequant(const block_q4_1 * x) const { + return b4.dequant(x->qs); + } +}; + +struct HBitDequantizer { + const __m256i shuffle = _mm256_set_epi64x(0x0303030303030303, 0x0202020202020202, 0x0101010101010101, 0x0000000000000000); + const __m256i mask = _mm256_set1_epi64x(0x7fbfdfeff7fbfdfe); + const __m256i minus1 = _mm256_set1_epi64x(-1); + inline __m256i to_bytes(const uint8_t * bits) const { + // Note: Data in all ggml quants is at least 2-byte aligned. + // => we can cast to uint16_t and use or on two consecutive entries + // which is faster than memcpy + const uint16_t * aux16 = (const uint16_t *)bits; + const uint32_t aux32 = aux16[0] | (aux16[1] << 16); + //uint32_t aux32; memcpy(&aux32, bits, sizeof(uint32_t)); + __m256i bytes = _mm256_shuffle_epi8(_mm256_set1_epi32(aux32), shuffle); + bytes = _mm256_or_si256(bytes, mask); + return _mm256_cmpeq_epi8(bytes, minus1); + } +}; + +struct Q5_0_Dequantizer { + Dequantizer4bit b4; + HBitDequantizer hbit; + const __m256i mh = _mm256_set1_epi8((char)0xF0); + inline __m256i dequant(const block_q5_0 * x) const { + const __m256i vqh = _mm256_andnot_si256(hbit.to_bytes(x->qh), mh); + return _mm256_or_si256(b4.dequant(x->qs), vqh); + } +}; + +template +struct Q5_1_Dequantizer { + Dequantizer4bit b4; + HBitDequantizer hbit; + const __m256i mh = _mm256_set1_epi8(0x10); + inline __m256i dequant(const Q5 * x) const { + const __m256i vqh = _mm256_and_si256(hbit.to_bytes(x->qh), mh); + return _mm256_or_si256(b4.dequant(x->qs), vqh); + } +}; +struct Q6_0_1_Dequantizer { + Dequantizer4bit b4; + const __m256i mh = _mm256_set1_epi8(0x30); + const __m256i shift1 = _mm256_set_epi64x(0, 2, 0, 4); + const __m256i shift2 = _mm256_set_epi64x(2, 0, 0, 0); + inline __m256i dequant(const block_q6_0 * x) const { + uint64_t aux64; std::memcpy(&aux64, x->qh, 8); + auto h256 = _mm256_sllv_epi64(_mm256_set1_epi64x(aux64), shift1); + return _mm256_or_si256(b4.dequant(x->qs), _mm256_and_si256(_mm256_srlv_epi64(h256, shift2), mh)); + } +}; +struct Q6_0_Dequantizer { + Q6_0_1_Dequantizer deq; + inline __m256i dequant(const block_q6_0 * x) const { + return _mm256_add_epi8(deq.dequant(x), _mm256_set1_epi8(-32)); + } +}; + +template +struct Q_Unpacker { + Q_Unpacker(const void * vx, size_t bx) : cx_0((const char *)vx), x((const Q*)cx_0), bx(bx) {} + + const char * cx_0; + const Q * x; + size_t bx; + + Scales scales; + Dequantizer deq; + + __m256i qx[4]; + + inline const __m256i* quants() const { return qx; } + + inline void set_row(int ix) { x = (const Q*)(cx_0 + ix*bx); } + + inline auto set_block_4(int i) { + for (int j = 0; j < 4; ++j) { + qx[j] = deq.dequant(x + 4*i + j); + } + return scales.prepare4(x + 4*i); + } + inline auto set_block(int i) { + qx[0] = deq.dequant(x + i); + return scales.prepare1(x + i); + } +}; + +struct Q8_0_Unpacker final : public Q_Unpacker { + Q8_0_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82S; + inline static int block_size() { return QK8_0; } +}; +struct Q8_0_1_Unpacker final : public Q_Unpacker, Q8_0_1_Dequantizer> { + Q8_0_1_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82; + inline static int block_size() { return QK8_0; } +}; +struct Q8_0_2_Unpacker final : public Q_Unpacker { + Q8_0_2_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82; + inline static int block_size() { return QK8_0; } +}; +struct Q4_0_Unpacker final : public Q_Unpacker { + Q4_0_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ80; + inline static int block_size() { return QK4_0; } +}; +struct Q4_0_1_Unpacker final : public Q_Unpacker, Q4_0_1_Dequantizer> { + Q4_0_1_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + //using Sum4T = Sum4TypeQ82; + using Sum4T = Sum4q4; + inline static int block_size() { return QK4_0; } +}; +struct MXFP4_Unpacker final : public Q_Unpacker, MXFP4_Dequantizer> { + MXFP4_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82; + inline static int block_size() { return QK4_NL; } +}; +struct IQ4_NL_UnpackerU final : public Q_Unpacker, IQ4_NL_DequantizerU> { + IQ4_NL_UnpackerU(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82; + inline static int block_size() { return QK4_NL; } +}; +struct IQ4_NL_UnpackerS final : public Q_Unpacker { + IQ4_NL_UnpackerS(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82S; + inline static int block_size() { return QK4_NL; } +}; +struct Q5_0_Unpacker final : public Q_Unpacker { + Q5_0_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ80; + inline static int block_size() { return QK5_0; } +}; +struct Q5_0_1_Unpacker final : public Q_Unpacker, Q5_1_Dequantizer> { + Q5_0_1_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82; + inline static int block_size() { return QK5_0; } +}; +struct Q4_1_Unpacker final : public Q_Unpacker { + Q4_1_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82; + inline static int block_size() { return QK4_1; } +}; +struct Q5_1_Unpacker final : public Q_Unpacker> { + Q5_1_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82; + inline static int block_size() { return QK5_1; } +}; +struct Q6_0_1_Unpacker final : public Q_Unpacker, Q6_0_1_Dequantizer> { + Q6_0_1_Unpacker(const void * vx, size_t bx) : Q_Unpacker(vx, bx) {} + using Sum4T = Sum4TypeQ82; + inline static int block_size() { return QK6_0; } +}; + +#ifdef HAVE_FANCY_SIMD +template +static void mul_mat_iq4_nl_r4_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + auto m4 = _mm512_set1_epi8(0xf); + auto values = load_iq4nl_values_512(); + int nb = n / QK4_NL; + __m512 acc[2*nrc_y] = {}; + __m512i qx[4]; + float d8[8*nrc_y]; + auto prepare = [&qx, &m4, &values] (const block_iq4_nl_r4& iq4l, const block_iq4_nl_r4& iq4h) { + auto scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq4l.d)); + auto scales1 = _mm256_set_m128(scales128, scales128); + scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq4h.d)); + auto scales2 = _mm256_set_m128(scales128, scales128); + auto scales = _mm512_insertf32x8(_mm512_castps256_ps512(scales1), scales2, 1); + auto bits1 = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq4l.qs+0)), + _mm256_loadu_si256((const __m256i *)iq4h.qs+0), 1); + auto bits2 = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq4l.qs+1)), + _mm256_loadu_si256((const __m256i *)iq4h.qs+1), 1); + qx[0] = _mm512_shuffle_epi8(values, _mm512_and_si512(bits1, m4)); + qx[1] = _mm512_shuffle_epi8(values, _mm512_and_si512(bits2, m4)); + qx[2] = _mm512_shuffle_epi8(values, _mm512_and_si512(_mm512_srli_epi16(bits1, 4), m4)); + qx[3] = _mm512_shuffle_epi8(values, _mm512_and_si512(_mm512_srli_epi16(bits2, 4), m4)); + return scales; + }; + auto dot = [&qx] (__m256i y8) { + auto y = _mm512_inserti32x8(_mm512_castsi256_si512(y8), y8, 1); + auto sumi = _mm512_setzero_si512(); + sumi = _mm512_dpbusd_epi32(sumi, qx[0], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + sumi = _mm512_dpbusd_epi32(sumi, qx[1], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + sumi = _mm512_dpbusd_epi32(sumi, qx[2], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + sumi = _mm512_dpbusd_epi32(sumi, qx[3], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + return sumi; + }; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_iq4_nl_r4 * iq4l = (const block_iq4_nl_r4 *)((const char *)vx + (ix+0)*bx); + const block_iq4_nl_r4 * iq4h = (const block_iq4_nl_r4 *)((const char *)vx + (ix+4)*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + _mm256_storeu_ps(d8+8*iy, convert_scales((const uint16_t *)q8.y[iy][ib4].d)); + } + for (int k = 0; k < 4; ++k) { + auto scales = prepare(iq4l[4*ib4+k], iq4h[4*ib4+k]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(_mm256_loadu_si256((const __m256i*)q8.y[iy][ib4].qs+k)); + auto dy = _mm512_set1_ps(d8[8*iy+k]); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(d8[8*iy+k+4]), acc[2*iy+1]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = prepare(iq4l[ib], iq4h[ib]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_1 *)q8.y[iy]; + auto sumi = dot(_mm256_loadu_si256((const __m256i*)qy[ib].qs)); + ggml_bf16_t d, s; d.bits = qy[ib].d; s.bits = qy[ib].s; + auto dy = _mm512_set1_ps(GGML_BF16_TO_FP32(d)); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(GGML_BF16_TO_FP32(s)), acc[2*iy+1]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum512 = _mm512_fmadd_ps(_mm512_set1_ps(-64.f), acc[2*iy+1], acc[2*iy+0]); + acc[2*iy+0] = acc[2*iy+1] = _mm512_setzero_ps(); + auto sum1 = _mm_add_ps(_mm512_extractf32x4_ps(sum512, 0), _mm512_extractf32x4_ps(sum512, 1)); + auto sum2 = _mm_add_ps(_mm512_extractf32x4_ps(sum512, 2), _mm512_extractf32x4_ps(sum512, 3)); + info.store(ix+0, iy, sum1); + info.store(ix+4, iy, sum2); + } + } +} +#else +template +static void mul_mat_iq4_nl_r4_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto m4 = _mm256_set1_epi8(0xf); +#ifndef HAVE_VNNI256 + auto m1 = _mm256_set1_epi16(1); +#endif + auto values128 = _mm_loadu_si128((const __m128i *)iq4k_values); + auto values = MM256_SET_M128I(values128, values128); + int nb = n / QK4_NL; + __m256 acc[nrc_y] = {}; + __m256i qs[4]; + float d8[4*nrc_y]; + auto prepare = [&qs, &values, &m4] (const block_iq4_nl_r4& iq4) { + auto scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq4.d)); + auto scales = _mm256_set_m128(scales128, scales128); + auto bits1 = _mm256_loadu_si256((const __m256i *)iq4.qs+0); + auto bits2 = _mm256_loadu_si256((const __m256i *)iq4.qs+1); + qs[0] = _mm256_shuffle_epi8(values, _mm256_and_si256(bits1, m4)); + qs[1] = _mm256_shuffle_epi8(values, _mm256_and_si256(bits2, m4)); + qs[2] = _mm256_shuffle_epi8(values, _mm256_and_si256(_mm256_srli_epi16(bits1, 4), m4)); + qs[3] = _mm256_shuffle_epi8(values, _mm256_and_si256(_mm256_srli_epi16(bits2, 4), m4)); + return scales; + }; +#ifdef HAVE_VNNI256 + auto dot = [&qs] (__m256i y) { + auto sumi = _mm256_setzero_si256(); + sumi = ggml_mm256_dpbusd_epi32(sumi, _mm256_sign_epi8(qs[0], qs[0]), _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qs[0])); + sumi = ggml_mm256_dpbusd_epi32(sumi, _mm256_sign_epi8(qs[1], qs[1]), _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qs[1])); + sumi = ggml_mm256_dpbusd_epi32(sumi, _mm256_sign_epi8(qs[2], qs[2]), _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qs[2])); + sumi = ggml_mm256_dpbusd_epi32(sumi, _mm256_sign_epi8(qs[3], qs[3]), _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qs[3])); + return sumi; + }; +#else + auto dot = [&qs, &m1] (__m256i y) { + auto u1 = _mm256_sign_epi8(qs[0], qs[0]); + auto u2 = _mm256_sign_epi8(qs[1], qs[1]); + auto sumi1 = _mm256_add_epi32( + _mm256_madd_epi16(m1, _mm256_maddubs_epi16(u1, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qs[0]))), + _mm256_madd_epi16(m1, _mm256_maddubs_epi16(u2, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qs[1])))); + u1 = _mm256_sign_epi8(qs[2], qs[2]); + u2 = _mm256_sign_epi8(qs[3], qs[3]); + auto sumi2 = _mm256_add_epi32( + _mm256_madd_epi16(m1, _mm256_maddubs_epi16(u1, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qs[2]))), + _mm256_madd_epi16(m1, _mm256_maddubs_epi16(u2, _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qs[3])))); + return _mm256_add_epi32(sumi1, sumi2); + }; +#endif + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_iq4_nl_r4 * iq4 = (const block_iq4_nl_r4 *)((const char *)vx + ix*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + auto aux = _mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)q8.y[iy][ib4].d)), 16); + _mm_storeu_ps(d8+4*iy, _mm_castsi128_ps(aux)); + } + for (int k = 0; k < 4; ++k) { + auto scales = prepare(iq4[4*ib4+k]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(_mm256_loadu_si256((const __m256i*)q8.y[iy][ib4].qs+k)); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[4*iy+k])); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = prepare(iq4[ib]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_1 *)q8.y[iy]; + auto sumi = dot(_mm256_loadu_si256((const __m256i*)qy[ib].qs)); + ggml_bf16_t d{qy[ib].d}; + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(GGML_BF16_TO_FP32(d))); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm_add_ps(_mm256_castps256_ps128(acc[iy]), _mm256_extractf128_ps(acc[iy], 1)); + info.store(ix, iy, sum); + acc[iy] = _mm256_setzero_ps(); + } + } +} +#endif + +inline void prepare_q4_0_quants_avx2(const uint8_t * qs, __m256i * v, const __m256i& m4) { + auto bits1 = _mm256_loadu_si256((const __m256i *)qs+0); + auto bits2 = _mm256_loadu_si256((const __m256i *)qs+1); + auto bits3 = _mm256_loadu_si256((const __m256i *)qs+2); + auto bits4 = _mm256_loadu_si256((const __m256i *)qs+3); + v[0] = _mm256_and_si256(bits1, m4); + v[1] = _mm256_and_si256(bits2, m4); + v[2] = _mm256_and_si256(bits3, m4); + v[3] = _mm256_and_si256(bits4, m4); + v[4] = _mm256_and_si256(_mm256_srli_epi16(bits1, 4), m4); + v[5] = _mm256_and_si256(_mm256_srli_epi16(bits2, 4), m4); + v[6] = _mm256_and_si256(_mm256_srli_epi16(bits3, 4), m4); + v[7] = _mm256_and_si256(_mm256_srli_epi16(bits4, 4), m4); +} + +inline __m256i accum_q4_0_quants(const __m256i * v, const int8_t * qs) { + auto y4l = _mm_loadu_si128((const __m128i*)qs+0); + auto y4h = _mm_loadu_si128((const __m128i*)qs+1); + auto yl = MM256_SET_M128I(y4l, y4l); + auto yh = MM256_SET_M128I(y4h, y4h); +#ifdef HAVE_FANCY_SIMD + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, v[0], _mm256_shuffle_epi32(yl, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, v[1], _mm256_shuffle_epi32(yl, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, v[2], _mm256_shuffle_epi32(yl, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, v[3], _mm256_shuffle_epi32(yl, 0xff)); + sumi = _mm256_dpbusd_epi32(sumi, v[4], _mm256_shuffle_epi32(yh, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, v[5], _mm256_shuffle_epi32(yh, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, v[6], _mm256_shuffle_epi32(yh, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, v[7], _mm256_shuffle_epi32(yh, 0xff)); +#else + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(v[0], _mm256_shuffle_epi32(yl, 0x00)), + _mm256_maddubs_epi16(v[1], _mm256_shuffle_epi32(yl, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(v[2], _mm256_shuffle_epi32(yl, 0xaa)), + _mm256_maddubs_epi16(v[3], _mm256_shuffle_epi32(yl, 0xff))); + auto sumi3 = _mm256_add_epi16(_mm256_maddubs_epi16(v[4], _mm256_shuffle_epi32(yh, 0x00)), + _mm256_maddubs_epi16(v[5], _mm256_shuffle_epi32(yh, 0x55))); + auto sumi4 = _mm256_add_epi16(_mm256_maddubs_epi16(v[6], _mm256_shuffle_epi32(yh, 0xaa)), + _mm256_maddubs_epi16(v[7], _mm256_shuffle_epi32(yh, 0xff))); + auto sumi = _mm256_madd_epi16(_mm256_set1_epi16(1), _mm256_add_epi16(_mm256_add_epi16(sumi1, sumi2), _mm256_add_epi16(sumi3, sumi4))); +#endif + return sumi; +} + +template +static void mul_mat_q4_0_r8_q8_2_avx2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + auto m4 = _mm256_set1_epi8(0xf); + int nb = n / QK4_NL; + __m256i v[8]; + if constexpr (nrc_y == 1) { + union { __m256 vec; float val[8]; } helper; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_iq4_nl_r8 * iq4 = (const block_iq4_nl_r8 *)((const char *)vx + ix*bx); + auto acc1 = _mm256_setzero_ps(); + auto acc2 = _mm256_setzero_ps(); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + helper.vec = convert_scales((const uint16_t *)q8.y[0][ib4].d); + for (int k = 0; k < 4; ++k) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4[4*ib4+k].d)); + prepare_q4_0_quants_avx2(iq4[4*ib4+k].qs, v, m4); + auto sumi = accum_q4_0_quants(v, q8.y[0][ib4].qs+32*k); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(helper.val[k])); + acc1 = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc1); + acc2 = _mm256_fmadd_ps(scales, _mm256_set1_ps(helper.val[k+4]), acc2); + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto qy = (const block_q8_2 *)q8.y[0]; + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4[ib].d)); + prepare_q4_0_quants_avx2(iq4[ib].qs, v, m4); + auto sumi = accum_q4_0_quants(v, qy[ib].qs); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8)); + acc1 = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc1); + acc2 = _mm256_fmadd_ps(scales, _mm256_set1_ps(m8), acc2); + } + acc1 = _mm256_fmadd_ps(acc2, _mm256_set1_ps(-8.f), acc1); + info.store(ix, 0, acc1); + } + } + else { + __m256 acc[nrc_y] = {}; + float d8[8*nrc_y]; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_iq4_nl_r8 * iq4 = (const block_iq4_nl_r8 *)((const char *)vx + ix*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + { + __m256 d4[4]; + for (int k = 0; k < 4; ++k) { + d4[k] = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4[4*ib4+k].d)); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto scales = convert_scales((const uint16_t *)q8.y[iy][ib4].d); + _mm256_storeu_ps(d8 + 8*iy, scales); + auto m4 = _mm256_extractf128_ps(scales, 1); + auto m8 = _mm256_set_m128(m4, m4); + auto sumf = _mm256_mul_ps(d4[0], _mm256_shuffle_ps(m8, m8, 0x00)); + sumf = _mm256_fmadd_ps(d4[1], _mm256_shuffle_ps(m8, m8, 0x55), sumf); + sumf = _mm256_fmadd_ps(d4[2], _mm256_shuffle_ps(m8, m8, 0xaa), sumf); + sumf = _mm256_fmadd_ps(d4[3], _mm256_shuffle_ps(m8, m8, 0xff), sumf); + acc[iy] = _mm256_fmadd_ps(sumf, _mm256_set1_ps(-8.f), acc[iy]); + } + } + for (int k = 0; k < 4; ++k) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4[4*ib4+k].d)); + prepare_q4_0_quants_avx2(iq4[4*ib4+k].qs, v, m4); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = accum_q4_0_quants(v, q8.y[iy][ib4].qs+32*k); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[8*iy+k])); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4[ib].d)); + auto scales_m = _mm256_mul_ps(scales, _mm256_set1_ps(-8.f)); + prepare_q4_0_quants_avx2(iq4[ib].qs, v, m4); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_2 *)q8.y[iy]; + auto sumi = accum_q4_0_quants(v, qy[ib].qs); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8)); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + acc[iy] = _mm256_fmadd_ps(scales_m, _mm256_set1_ps(m8), acc[iy]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = _mm256_setzero_ps(); + } + } + } +} + +#ifdef HAVE_FANCY_SIMD +template +static void mul_mat_q4_0_r8_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + if constexpr (nrc_y == 1) { + mul_mat_q4_0_r8_q8_2_avx2<1>(n, vx, bx, info, nrc_x); + return; + } + GGML_ASSERT(nrc_x%16 == 0); + Q8 q8(info); + auto m4 = _mm512_set1_epi8(0xf); + int nb = n / QK4_NL; + __m512 acc[2*nrc_y] = {}; + __m512i qx[8]; + auto prepare = [&qx, &m4] (const block_iq4_nl_r8& iq4l, const block_iq4_nl_r8& iq4h) { + auto scales1 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4l.d)); + auto scales2 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq4h.d)); + auto scales = _mm512_insertf32x8(_mm512_castps256_ps512(scales1), scales2, 1); + for (int j = 0; j < 4; ++j) { + auto bits = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq4l.qs+j)), + _mm256_loadu_si256((const __m256i *)iq4h.qs+j), 1); + qx[j+0] = _mm512_and_si512(bits, m4); + qx[j+4] = _mm512_and_si512(_mm512_srli_epi16(bits, 4), m4); + } + return scales; + }; + auto dot = [&qx] (const int8_t * qy) { + auto y4l = _mm_loadu_si128((const __m128i*)qy+0); + auto y4h = _mm_loadu_si128((const __m128i*)qy+1); + auto y8l = MM256_SET_M128I(y4l, y4l); + auto y8h = MM256_SET_M128I(y4h, y4h); + auto yl = _mm512_inserti32x8(_mm512_castsi256_si512(y8l), y8l, 1); + auto yh = _mm512_inserti32x8(_mm512_castsi256_si512(y8h), y8h, 1); + auto sumi = _mm512_setzero_si512(); + sumi = _mm512_dpbusd_epi32(sumi, qx[0], _mm512_shuffle_epi32(yl, _MM_PERM_ENUM(0x00))); + sumi = _mm512_dpbusd_epi32(sumi, qx[1], _mm512_shuffle_epi32(yl, _MM_PERM_ENUM(0x55))); + sumi = _mm512_dpbusd_epi32(sumi, qx[2], _mm512_shuffle_epi32(yl, _MM_PERM_ENUM(0xaa))); + sumi = _mm512_dpbusd_epi32(sumi, qx[3], _mm512_shuffle_epi32(yl, _MM_PERM_ENUM(0xff))); + sumi = _mm512_dpbusd_epi32(sumi, qx[4], _mm512_shuffle_epi32(yh, _MM_PERM_ENUM(0x00))); + sumi = _mm512_dpbusd_epi32(sumi, qx[5], _mm512_shuffle_epi32(yh, _MM_PERM_ENUM(0x55))); + sumi = _mm512_dpbusd_epi32(sumi, qx[6], _mm512_shuffle_epi32(yh, _MM_PERM_ENUM(0xaa))); + sumi = _mm512_dpbusd_epi32(sumi, qx[7], _mm512_shuffle_epi32(yh, _MM_PERM_ENUM(0xff))); + return sumi; + }; + float d8[8*nrc_y]; + for (int ix = 0; ix < nrc_x; ix += 16) { + const block_iq4_nl_r8 * iq4l = (const block_iq4_nl_r8 *)((const char *)vx + (ix+0)*bx); + const block_iq4_nl_r8 * iq4h = (const block_iq4_nl_r8 *)((const char *)vx + (ix+8)*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + _mm256_storeu_ps(d8+8*iy, convert_scales((const uint16_t *)q8.y[iy][ib4].d)); + } + for (int k = 0; k < 4; ++k) { + auto scales = prepare(iq4l[4*ib4+k], iq4h[4*ib4+k]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(q8.y[iy][ib4].qs+32*k); + auto dy = _mm512_set1_ps(d8[8*iy+k]); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(d8[8*iy+k+4]), acc[2*iy+1]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = prepare(iq4l[ib], iq4h[ib]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_1 *)q8.y[iy]; + auto sumi = dot(qy[ib].qs); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto dy = _mm512_set1_ps(d8); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(m8), acc[2*iy+1]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm512_fmadd_ps(_mm512_set1_ps(-8.f), acc[2*iy+1], acc[2*iy+0]); + acc[2*iy+0] = acc[2*iy+1] = _mm512_setzero_ps(); + info.store(ix, iy, sum); + } + } +} +#else +template +static void mul_mat_q4_0_r8_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + mul_mat_q4_0_r8_q8_2_avx2(n, vx, bx, info, nrc_x); +} +#endif + +template +static void mul_mat_q5_0_r4_q8_2_avx2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto m4 = _mm256_set1_epi8(0xf); + auto m5 = _mm256_set1_epi8(0x10); +#ifndef HAVE_FANCY_SIMD + auto m1 = _mm256_set1_epi16(1); +#endif + auto mscale = _mm256_set_m128(_mm_set1_ps(-8.f), _mm_set1_ps(1.f)); + int nb = n / QK5_0; + __m256 acc[nrc_y] = {}; + __m256i qx[4]; + float d8[8*nrc_y]; + auto prepare = [&qx, &m4, &m5] (const block_q5_0_r4& iq5) { + auto scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq5.d)); + auto scales = _mm256_set_m128(scales128, scales128); + auto bits1 = _mm256_loadu_si256((const __m256i *)iq5.qs+0); + auto bits2 = _mm256_loadu_si256((const __m256i *)iq5.qs+1); + auto hbits = _mm_loadu_si128((const __m128i *)iq5.qh); + auto hb = MM256_SET_M128I(_mm_srli_epi16(hbits, 1), hbits); + qx[0] = _mm256_or_si256(_mm256_and_si256(bits1, m4), _mm256_and_si256(_mm256_slli_epi16(hb, 4), m5)); + qx[1] = _mm256_or_si256(_mm256_and_si256(bits2, m4), _mm256_and_si256(_mm256_slli_epi16(hb, 2), m5)); + qx[2] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(bits1, 4), m4), _mm256_and_si256(hb, m5)); + qx[3] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(bits2, 4), m4), _mm256_and_si256(_mm256_srli_epi16(hb, 2), m5));; + return scales; + }; +#ifdef HAVE_FANCY_SIMD + auto dot = [&qx] (__m256i y) { + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + return sumi; + }; +#else + auto dot = [&qx, &m1] (__m256i y) { + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[0], _mm256_shuffle_epi32(y, 0x00)), + _mm256_maddubs_epi16(qx[1], _mm256_shuffle_epi32(y, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[2], _mm256_shuffle_epi32(y, 0xaa)), + _mm256_maddubs_epi16(qx[3], _mm256_shuffle_epi32(y, 0xff))); + auto sumi = _mm256_madd_epi16(m1, _mm256_add_epi16(sumi1, sumi2)); + return sumi; + }; +#endif + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q5_0_r4 * iq5 = (const block_q5_0_r4 *)((const char *)vx + ix*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + auto scales = convert_scales((const uint16_t *)q8.y[iy][ib4].d); + _mm256_storeu_ps(d8 + 8*iy, _mm256_mul_ps(mscale, scales)); + } + for (int k = 0; k < 4; ++k) { + auto scales = prepare(iq5[4*ib4+k]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(_mm256_loadu_si256((const __m256i*)q8.y[iy][ib4].qs+k)); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[8*iy+k])); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + acc[iy] = _mm256_fmadd_ps(scales, _mm256_set1_ps(d8[8*iy+k+4]), acc[iy]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = prepare(iq5[ib]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_2 *)q8.y[iy]; + auto sumi = dot(_mm256_loadu_si256((const __m256i*)qy[ib].qs)); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8)); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + acc[iy] = _mm256_fmadd_ps(scales, _mm256_set1_ps(-8.f*m8), acc[iy]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm_add_ps(_mm256_castps256_ps128(acc[iy]), _mm256_extractf128_ps(acc[iy], 1)); + info.store(ix, iy, sum); + acc[iy] = _mm256_setzero_ps(); + } + } +} + +#ifdef HAVE_FANCY_SIMD +template +static void mul_mat_q5_0_r4_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + if constexpr (nrc_y == 1) { + mul_mat_q5_0_r4_q8_2_avx2<1>(n, vx, bx, info, nrc_x); + } else { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + auto m4 = _mm512_set1_epi8(0xf); + auto m5 = _mm512_set1_epi8(0x10); + int nb = n / QK5_0; + __m512 acc[2*nrc_y] = {}; + __m512i qx[4]; + float d8[8*nrc_y]; + auto prepare = [&qx, &m4, &m5] (const block_q5_0_r4& iq5l, const block_q5_0_r4& iq5h) { + auto scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq5l.d)); + auto scales1 = _mm256_set_m128(scales128, scales128); + scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq5h.d)); + auto scales2 = _mm256_set_m128(scales128, scales128); + auto scales = _mm512_insertf32x8(_mm512_castps256_ps512(scales1), scales2, 1); + auto bits1 = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq5l.qs+0)), + _mm256_loadu_si256((const __m256i *)iq5h.qs+0), 1); + auto bits2 = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq5l.qs+1)), + _mm256_loadu_si256((const __m256i *)iq5h.qs+1), 1); + auto hbits1 = _mm_loadu_si128((const __m128i *)iq5l.qh); + auto hbits2 = _mm_loadu_si128((const __m128i *)iq5h.qh); + auto hb1 = MM256_SET_M128I(_mm_srli_epi16(hbits1, 1), hbits1); + auto hb2 = MM256_SET_M128I(_mm_srli_epi16(hbits2, 1), hbits2); + auto hb = _mm512_inserti32x8(_mm512_castsi256_si512(hb1), hb2, 1); + qx[0] = _mm512_or_si512(_mm512_and_si512(bits1, m4), _mm512_and_si512(_mm512_slli_epi16(hb, 4), m5)); + qx[1] = _mm512_or_si512(_mm512_and_si512(bits2, m4), _mm512_and_si512(_mm512_slli_epi16(hb, 2), m5)); + qx[2] = _mm512_or_si512(_mm512_and_si512(_mm512_srli_epi16(bits1, 4), m4), _mm512_and_si512(hb, m5)); + qx[3] = _mm512_or_si512(_mm512_and_si512(_mm512_srli_epi16(bits2, 4), m4), _mm512_and_si512(_mm512_srli_epi16(hb, 2), m5)); + return scales; + }; + auto dot = [&qx] (__m256i y8) { + auto y = _mm512_inserti32x8(_mm512_castsi256_si512(y8), y8, 1); + auto sumi = _mm512_setzero_si512(); + sumi = _mm512_dpbusd_epi32(sumi, qx[0], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + sumi = _mm512_dpbusd_epi32(sumi, qx[1], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + sumi = _mm512_dpbusd_epi32(sumi, qx[2], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + sumi = _mm512_dpbusd_epi32(sumi, qx[3], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + return sumi; + }; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q5_0_r4 * iq5l = (const block_q5_0_r4 *)((const char *)vx + (ix+0)*bx); + const block_q5_0_r4 * iq5h = (const block_q5_0_r4 *)((const char *)vx + (ix+4)*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + _mm256_storeu_ps(d8+8*iy, convert_scales((const uint16_t *)q8.y[iy][ib4].d)); + } + for (int k = 0; k < 4; ++k) { + auto scales = prepare(iq5l[4*ib4+k], iq5h[4*ib4+k]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(_mm256_loadu_si256((const __m256i*)q8.y[iy][ib4].qs+k)); + auto dy = _mm512_set1_ps(d8[8*iy+k]); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(d8[8*iy+k+4]), acc[2*iy+1]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = prepare(iq5l[ib], iq5h[ib]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_2 *)q8.y[iy]; + auto sumi = dot(_mm256_loadu_si256((const __m256i*)qy[ib].qs)); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto dy = _mm512_set1_ps(d8); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(m8), acc[2*iy+1]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum512 = _mm512_fmadd_ps(_mm512_set1_ps(-8.f), acc[2*iy+1], acc[2*iy+0]); + acc[2*iy+0] = acc[2*iy+1] = _mm512_setzero_ps(); + auto sum1 = _mm_add_ps(_mm512_extractf32x4_ps(sum512, 0), _mm512_extractf32x4_ps(sum512, 1)); + auto sum2 = _mm_add_ps(_mm512_extractf32x4_ps(sum512, 2), _mm512_extractf32x4_ps(sum512, 3)); + info.store(ix+0, iy, sum1); + info.store(ix+4, iy, sum2); + } + } + } +} +#else +template +static void mul_mat_q5_0_r4_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + mul_mat_q5_0_r4_q8_2_avx2(n, vx, bx, info, nrc_x); +} +#endif + +template +static void mul_mat_q6_0_r4_q8_2_avx2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + auto m4 = _mm256_set1_epi8(0xf); + auto m6 = _mm256_set1_epi8(0x30); + auto mscale = _mm256_set_m128(_mm_set1_ps(-16.f), _mm_set1_ps(1.f)); +#ifndef HAVE_FANCY_SIMD + auto m1 = _mm256_set1_epi16(1); +#endif + int nb = n / QK6_0; + __m256 acc[nrc_y] = {}; + float d8[8*nrc_y]; + __m256i qx[4]; + auto prepare = [&qx, &m4, &m6] (const block_q6_0_r4& iq6) { + auto scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq6.d)); + auto scales = _mm256_set_m128(scales128, scales128); + auto bits1 = _mm256_loadu_si256((const __m256i *)iq6.qs+0); + auto bits2 = _mm256_loadu_si256((const __m256i *)iq6.qs+1); + auto hbits = _mm256_loadu_si256((const __m256i *)iq6.qh); + qx[0] = _mm256_or_si256(_mm256_and_si256(bits1, m4), _mm256_and_si256(_mm256_slli_epi16(hbits, 4), m6)); + qx[1] = _mm256_or_si256(_mm256_and_si256(bits2, m4), _mm256_and_si256(_mm256_slli_epi16(hbits, 2), m6)); + qx[2] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(bits1, 4), m4), _mm256_and_si256(hbits, m6)); + qx[3] = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(bits2, 4), m4), _mm256_and_si256(_mm256_srli_epi16(hbits, 2), m6)); + return scales; + }; +#ifdef HAVE_FANCY_SIMD + auto dot = [&qx] (__m256i y) { + auto sumi = _mm256_dpbusd_epi32(_mm256_setzero_si256(), qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + return sumi; + }; +#else + auto dot = [&qx, &m1] (__m256i y) { + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[0], _mm256_shuffle_epi32(y, 0x00)), + _mm256_maddubs_epi16(qx[1], _mm256_shuffle_epi32(y, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[2], _mm256_shuffle_epi32(y, 0xaa)), + _mm256_maddubs_epi16(qx[3], _mm256_shuffle_epi32(y, 0xff))); + auto sumi = _mm256_add_epi32(_mm256_madd_epi16(m1, sumi1), _mm256_madd_epi16(m1, sumi2)); + return sumi; + }; +#endif + for (int ix = 0; ix < nrc_x; ix += 4) { + const block_q6_0_r4 * iq6 = (const block_q6_0_r4 *)((const char *)vx + ix*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + auto scales = convert_scales((const uint16_t *)q8.y[iy][ib4].d); + _mm256_storeu_ps(d8 + 8*iy, _mm256_mul_ps(scales, mscale)); + } + for (int k = 0; k < 4; ++k) { + auto scales = prepare(iq6[4*ib4+k]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(_mm256_loadu_si256((const __m256i*)q8.y[iy][ib4].qs+k)); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[8*iy+k])); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + acc[iy] = _mm256_fmadd_ps(scales, _mm256_set1_ps(d8[8*iy+k+4]), acc[iy]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = prepare(iq6[ib]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_2 *)q8.y[iy]; + auto sumi = dot(_mm256_loadu_si256((const __m256i*)qy[ib].qs)); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8)); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + acc[iy] = _mm256_fmadd_ps(scales, _mm256_set1_ps(-16.f*m8), acc[iy]); + } + } + + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum = _mm_add_ps(_mm256_castps256_ps128(acc[iy]), _mm256_extractf128_ps(acc[iy], 1)); + info.store(ix, iy, sum); + acc[iy] = _mm256_setzero_ps(); + } + } +} + +#ifdef HAVE_FANCY_SIMD +template +static void mul_mat_q6_0_r4_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + if constexpr (nrc_y == 1) { + mul_mat_q6_0_r4_q8_2_avx2<1>(n, vx, bx, info, nrc_x); + } else { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + auto m4 = _mm512_set1_epi8(0xf); + auto m6 = _mm512_set1_epi8(0x30); + int nb = n / QK6_0; + __m512 acc[2*nrc_y] = {}; + __m512i qx[4]; + float d8[8*nrc_y]; + auto prepare = [&qx, &m4, &m6] (const block_q6_0_r4& iq6l, const block_q6_0_r4& iq6h) { + auto scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq6l.d)); + auto scales1 = _mm256_set_m128(scales128, scales128); + scales128 = _mm_cvtph_ps(_mm_loadl_epi64((const __m128i *)iq6h.d)); + auto scales2 = _mm256_set_m128(scales128, scales128); + auto scales = _mm512_insertf32x8(_mm512_castps256_ps512(scales1), scales2, 1); + auto bits1 = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq6l.qs+0)), + _mm256_loadu_si256((const __m256i *)iq6h.qs+0), 1); + auto bits2 = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)iq6l.qs+1)), + _mm256_loadu_si256((const __m256i *)iq6h.qs+1), 1); + auto hbits1 = _mm256_loadu_si256((const __m256i *)iq6l.qh); + auto hbits2 = _mm256_loadu_si256((const __m256i *)iq6h.qh); + auto hb = _mm512_inserti32x8(_mm512_castsi256_si512(hbits1), hbits2, 1); + qx[0] = _mm512_and_si512(bits1, m4) | _mm512_and_si512(_mm512_slli_epi16(hb, 4), m6); + qx[1] = _mm512_and_si512(bits2, m4) | _mm512_and_si512(_mm512_slli_epi16(hb, 2), m6);; + qx[2] = _mm512_and_si512(_mm512_srli_epi16(bits1, 4), m4) | _mm512_and_si512(hb, m6); + qx[3] = _mm512_and_si512(_mm512_srli_epi16(bits2, 4), m4) | _mm512_and_si512(_mm512_srli_epi16(hb, 2), m6); + return scales; + }; + auto dot = [&qx] (__m256i y8) { + auto y = _mm512_inserti32x8(_mm512_castsi256_si512(y8), y8, 1); + auto sumi = _mm512_setzero_si512(); + sumi = _mm512_dpbusd_epi32(sumi, qx[0], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x00))); + sumi = _mm512_dpbusd_epi32(sumi, qx[1], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0x55))); + sumi = _mm512_dpbusd_epi32(sumi, qx[2], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xaa))); + sumi = _mm512_dpbusd_epi32(sumi, qx[3], _mm512_shuffle_epi32(y, _MM_PERM_ENUM(0xff))); + return sumi; + }; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q6_0_r4 * iq6l = (const block_q6_0_r4 *)((const char *)vx + (ix+0)*bx); + const block_q6_0_r4 * iq6h = (const block_q6_0_r4 *)((const char *)vx + (ix+4)*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + auto scales = convert_scales((const uint16_t *)q8.y[iy][ib4].d); + _mm256_storeu_ps(d8 + 8*iy, scales); + } + for (int k = 0; k < 4; ++k) { + auto scales = prepare(iq6l[4*ib4+k], iq6h[4*ib4+k]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(_mm256_loadu_si256((const __m256i*)q8.y[iy][ib4].qs+k)); + auto dy = _mm512_set1_ps(d8[8*iy+k]); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(d8[8*iy+k+4]), acc[2*iy+1]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = prepare(iq6l[ib], iq6h[ib]); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_2 *)q8.y[iy]; + auto sumi = dot(_mm256_loadu_si256((const __m256i*)qy[ib].qs)); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto dy = _mm512_set1_ps(d8); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(m8), acc[2*iy+1]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum512 = _mm512_fmadd_ps(_mm512_set1_ps(-16.f), acc[2*iy+1], acc[2*iy+0]); + acc[2*iy+0] = acc[2*iy+1] = _mm512_setzero_ps(); + auto sum1 = _mm_add_ps(_mm512_extractf32x4_ps(sum512, 0), _mm512_extractf32x4_ps(sum512, 1)); + auto sum2 = _mm_add_ps(_mm512_extractf32x4_ps(sum512, 2), _mm512_extractf32x4_ps(sum512, 3)); + info.store(ix+0, iy, sum1); + info.store(ix+4, iy, sum2); + } + } + } +} +#else +template +static void mul_mat_q6_0_r4_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + mul_mat_q6_0_r4_q8_2_avx2(n, vx, bx, info, nrc_x); +} +#endif + +#ifdef HAVE_FANCY_SIMD +inline __m512i qx_r8_q8_dot_product(const __m512i * qx, const int8_t * y) { + auto y4l = _mm_loadu_si128((const __m128i*)y+0); + auto y4h = _mm_loadu_si128((const __m128i*)y+1); + auto y8l = MM256_SET_M128I(y4l, y4l); + auto y8h = MM256_SET_M128I(y4h, y4h); + auto yl = _mm512_inserti32x8(_mm512_castsi256_si512(y8l), y8l, 1); + auto yh = _mm512_inserti32x8(_mm512_castsi256_si512(y8h), y8h, 1); + auto sumi = _mm512_setzero_si512(); + sumi = _mm512_dpbusd_epi32(sumi, qx[0], _mm512_shuffle_epi32(yl, _MM_PERM_ENUM(0x00))); + sumi = _mm512_dpbusd_epi32(sumi, qx[1], _mm512_shuffle_epi32(yl, _MM_PERM_ENUM(0x55))); + sumi = _mm512_dpbusd_epi32(sumi, qx[2], _mm512_shuffle_epi32(yl, _MM_PERM_ENUM(0xaa))); + sumi = _mm512_dpbusd_epi32(sumi, qx[3], _mm512_shuffle_epi32(yl, _MM_PERM_ENUM(0xff))); + sumi = _mm512_dpbusd_epi32(sumi, qx[4], _mm512_shuffle_epi32(yh, _MM_PERM_ENUM(0x00))); + sumi = _mm512_dpbusd_epi32(sumi, qx[5], _mm512_shuffle_epi32(yh, _MM_PERM_ENUM(0x55))); + sumi = _mm512_dpbusd_epi32(sumi, qx[6], _mm512_shuffle_epi32(yh, _MM_PERM_ENUM(0xaa))); + sumi = _mm512_dpbusd_epi32(sumi, qx[7], _mm512_shuffle_epi32(yh, _MM_PERM_ENUM(0xff))); + return sumi; +} +inline __m256i qx_r8_q8_dot_product(const __m256i * qx, const int8_t * y) { + auto y4l = _mm_loadu_si128((const __m128i*)y+0); + auto y4h = _mm_loadu_si128((const __m128i*)y+1); + auto yl = MM256_SET_M128I(y4l, y4l); + auto yh = MM256_SET_M128I(y4h, y4h); + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(yl, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(yl, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(yl, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(yl, 0xff)); + sumi = _mm256_dpbusd_epi32(sumi, qx[4], _mm256_shuffle_epi32(yh, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[5], _mm256_shuffle_epi32(yh, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[6], _mm256_shuffle_epi32(yh, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[7], _mm256_shuffle_epi32(yh, 0xff)); + return sumi; +} +inline __m256i q8_0_r8_dot_product(const uint8_t * x, const int8_t * y, __m256i * qx) { + for (int i = 0; i < 8; ++i) { + qx[i] = _mm256_add_epi8(_mm256_loadu_si256((const __m256i *)x+i), _mm256_set1_epi8(127)); + } + return qx_r8_q8_dot_product(qx, y); +} +template +static void mul_mat_q8_0_r8_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%16 == 0); + Q8 q8(info); + int nb = n / QK8_0; + if constexpr (nrc_y == 1) { + __m256 acc[2] = {}; + __m256i qx[8]; + float d8[8]; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q8_0_r8 * iq8 = (const block_q8_0_r8 *)((const char *)vx + ix*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + _mm256_storeu_ps(d8, convert_scales((const uint16_t *)q8.y[0][ib4].d)); + for (int k = 0; k < 4; ++k) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[4*ib4+k].d)); + auto sumi = q8_0_r8_dot_product((const uint8_t *)iq8[4*ib4+k].qs, q8.y[0][ib4].qs+32*k, qx); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[k])); + acc[0] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[0]); + acc[1] = _mm256_fmadd_ps(scales, _mm256_set1_ps(d8[k+4]), acc[1]); + } + } + if (4*(nb/4) < nb) { + auto qy = (const block_q8_2 *)q8.y[0]; + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[ib].d)); + auto sumi = q8_0_r8_dot_product((const uint8_t *)iq8[ib].qs, qy[ib].qs, qx); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8)); + acc[0] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[0]); + acc[1] = _mm256_fmadd_ps(scales, _mm256_set1_ps(m8), acc[1]); + } + } + info.store(ix, 0, _mm256_fmadd_ps(_mm256_set1_ps(-127.f), acc[1], acc[0])); + acc[0] = acc[1] = _mm256_setzero_ps(); + } + } else { + __m512 acc[2*nrc_y] = {}; + __m512i qx[8]; + float d8[8*nrc_y]; + for (int ix = 0; ix < nrc_x; ix += 16) { + const block_q8_0_r8 * q8l = (const block_q8_0_r8 *)((const char *)vx + (ix+0)*bx); + const block_q8_0_r8 * q8h = (const block_q8_0_r8 *)((const char *)vx + (ix+8)*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + _mm256_storeu_ps(d8+8*iy, convert_scales((const uint16_t *)q8.y[iy][ib4].d)); + } + for (int k = 0; k < 4; ++k) { + auto scales1 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8l[4*ib4+k].d)); + auto scales2 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8h[4*ib4+k].d)); + auto scales = _mm512_insertf32x8(_mm512_castps256_ps512(scales1), scales2, 1); + for (int j = 0; j < 8; ++j) { + qx[j] = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)q8l[4*ib4+k].qs+j)), + _mm256_loadu_si256((const __m256i *)q8h[4*ib4+k].qs+j), 1); + qx[j] = _mm512_add_epi8(qx[j], _mm512_set1_epi8(127)); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = qx_r8_q8_dot_product(qx, q8.y[iy][ib4].qs+32*k); + auto dy = _mm512_set1_ps(d8[8*iy+k]); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(d8[8*iy+k+4]), acc[2*iy+1]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales1 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8l[ib].d)); + auto scales2 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8h[ib].d)); + auto scales = _mm512_insertf32x8(_mm512_castps256_ps512(scales1), scales2, 1); + for (int j = 0; j < 8; ++j) { + qx[j] = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)q8l[ib].qs+j)), + _mm256_loadu_si256((const __m256i *)q8h[ib].qs+j), 1); + qx[j] = _mm512_add_epi8(qx[j], _mm512_set1_epi8(127)); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_2 *)q8.y[iy]; + auto sumi = qx_r8_q8_dot_product(qx, qy[ib].qs); + auto [d8, m8] = ScaleHelperQ8_2::prepare1(qy + ib); + auto dy = _mm512_set1_ps(d8); + acc[2*iy+0] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[2*iy+0]); + acc[2*iy+1] = _mm512_fmadd_ps(scales, _mm512_set1_ps(m8), acc[2*iy+1]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sum512 = _mm512_fmadd_ps(_mm512_set1_ps(-127.f), acc[2*iy+1], acc[2*iy+0]); + info.store(ix, iy, sum512); + acc[2*iy+0] = acc[2*iy+1] = _mm512_setzero_ps(); + } + } + } +} +#else +template +static void mul_mat_q8_0_r8_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + auto m1 = _mm256_set1_epi16(1); + int nb = n / QK8_0; + __m256 acc[nrc_y] = {}; + float d8[4*nrc_y]; + __m256i qx[4], sx[4]; + auto dot = [&qx, &sx, &m1] (const int8_t * qy) { + auto y128 = _mm_loadu_si128((const __m128i*)qy); + auto y = MM256_SET_M128I(y128, y128); +#ifdef HAVE_VNNI256 + auto sumi = _mm256_setzero_si256(); + sumi = ggml_mm256_dpbusd_epi32(sumi, sx[0], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0])); + sumi = ggml_mm256_dpbusd_epi32(sumi, sx[1], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1])); + sumi = ggml_mm256_dpbusd_epi32(sumi, sx[2], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2])); + sumi = ggml_mm256_dpbusd_epi32(sumi, sx[3], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qx[3])); + return sumi; +#else + auto sumi1 = _mm256_add_epi32( + _mm256_madd_epi16(m1, _mm256_maddubs_epi16(sx[0], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x00), qx[0]))), + _mm256_madd_epi16(m1, _mm256_maddubs_epi16(sx[1], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0x55), qx[1]))) + ); + auto sumi2 = _mm256_add_epi32( + _mm256_madd_epi16(m1, _mm256_maddubs_epi16(sx[2], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xaa), qx[2]))), + _mm256_madd_epi16(m1, _mm256_maddubs_epi16(sx[3], _mm256_sign_epi8(_mm256_shuffle_epi32(y, 0xff), qx[3]))) + ); + return _mm256_add_epi32(sumi1, sumi2); +#endif + }; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q8_0_r8 * iq8 = (const block_q8_0_r8 *)((const char *)vx + ix*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + auto scales = _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)q8.y[iy][ib4].d)), 16)); + _mm_storeu_ps(d8 + 4*iy, scales); + } + for (int k = 0; k < 4; ++k) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[4*ib4+k].d)); + __m256i sumi_first[nrc_y]; + for (int j = 0; j < 4; ++j) { + qx[j] = _mm256_loadu_si256((const __m256i *)iq8[4*ib4+k].qs+j); + sx[j] = _mm256_sign_epi8(qx[j], qx[j]); + } + for (int iy = 0; iy < nrc_y; ++iy) { + sumi_first[iy] = dot(q8.y[iy][ib4].qs+32*k); + } + for (int j = 0; j < 4; ++j) { + qx[j] = _mm256_loadu_si256((const __m256i *)iq8[4*ib4+k].qs+4+j); + sx[j] = _mm256_sign_epi8(qx[j], qx[j]); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = _mm256_add_epi32(sumi_first[iy], dot(q8.y[iy][ib4].qs+32*k+16)); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[4*iy+k])); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[ib].d)); + __m256i sumi_first[nrc_y]; + for (int j = 0; j < 4; ++j) { + qx[j] = _mm256_loadu_si256((const __m256i *)iq8[ib].qs+j); + sx[j] = _mm256_sign_epi8(qx[j], qx[j]); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_2 *)q8.y[iy]; + sumi_first[iy] = dot(qy[ib].qs); + } + for (int j = 0; j < 4; ++j) { + qx[j] = _mm256_loadu_si256((const __m256i *)iq8[ib].qs+4+j); + sx[j] = _mm256_sign_epi8(qx[j], qx[j]); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_2 *)q8.y[iy]; + auto sumi = _mm256_add_epi32(sumi_first[iy], dot(qy[ib].qs+16)); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(GGML_BF16_TO_FP32(ggml_bf16_t{qy[ib].d}))); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = _mm256_setzero_ps(); + } + } +} +#endif + +typedef struct { + ggml_half d[16]; + uint8_t qs[256]; +} block_q8_1_r8; + +#ifdef HAVE_FANCY_SIMD +template +static void mul_mat_q8_1_r8_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%16 == 0); + Q8 q8(info); + int nb = n / QK8_0; + if constexpr (nrc_y == 1) { + __m256 acc[1] = {}; + float d8[4]; + __m256i qx[4]; + auto dot = [&qx] (const int8_t * qy) { + auto y128 = _mm_loadu_si128((const __m128i*)qy); + auto y = MM256_SET_M128I(y128, y128); + auto sumi = _mm256_setzero_si256(); + sumi = _mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = _mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = _mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = _mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + return sumi; + }; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q8_1_r8 * iq8 = (const block_q8_1_r8 *)((const char *)vx + ix*bx); + for (int i4 = 0; i4 < nb/4; ++i4) { + { + __m256 mx[4]; + for (int ib32 = 0; ib32 < 4; ++ib32) mx[ib32] = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[4*i4+ib32].d+1)); + auto scales = _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)q8.y[0][i4].d)), 16)); + _mm_storeu_ps(d8, scales); + auto bsums4 = _mm_cvtepi32_ps(_mm_cvtepi16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[0][i4].d+4)))); + bsums4 = _mm_mul_ps(bsums4, scales); + auto bsums = _mm256_set_m128(bsums4, bsums4); + acc[0] = _mm256_fmadd_ps(mx[0], _mm256_shuffle_ps(bsums, bsums, 0x00), acc[0]); + acc[0] = _mm256_fmadd_ps(mx[1], _mm256_shuffle_ps(bsums, bsums, 0x55), acc[0]); + acc[0] = _mm256_fmadd_ps(mx[2], _mm256_shuffle_ps(bsums, bsums, 0xaa), acc[0]); + acc[0] = _mm256_fmadd_ps(mx[3], _mm256_shuffle_ps(bsums, bsums, 0xff), acc[0]); + } + for (int ib32 = 0; ib32 < 4; ++ib32) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[4*i4+ib32].d)); + for (int j = 0; j < 4; ++j) { + qx[j] = _mm256_loadu_si256((const __m256i *)iq8[4*i4+ib32].qs+j); + } + auto sumi = dot(q8.y[0][i4].qs+32*ib32); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[ib32])); + acc[0] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[0]); + for (int j = 0; j < 4; ++j) { + qx[j] = _mm256_loadu_si256((const __m256i *)iq8[4*i4+ib32].qs+4+j); + } + sumi = dot(q8.y[0][i4].qs+32*ib32+16); + d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[ib32])); + acc[0] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[0]); + } + } + info.store(ix, 0, acc[0]); + acc[0] = _mm256_setzero_ps(); + } + } else { + __m512 acc[nrc_y] = {}; + __m512i qx[8]; + float d8[4*nrc_y]; + for (int ix = 0; ix < nrc_x; ix += 16) { + const block_q8_1_r8 * q8l = (const block_q8_1_r8 *)((const char *)vx + (ix+0)*bx); + const block_q8_1_r8 * q8h = (const block_q8_1_r8 *)((const char *)vx + (ix+8)*bx); + for (int i4 = 0; i4 < nb/4; ++i4) { + { + __m512 mx[4]; + for (int ib32 = 0; ib32 < 4; ++ib32) { + auto mx_l = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8l[4*i4+ib32].d+1)); + auto mx_h = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8h[4*i4+ib32].d+1)); + mx[ib32] = _mm512_insertf32x8(_mm512_castps256_ps512(mx_l), mx_h, 1); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto scales128 = _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)q8.y[iy][i4].d)), 16)); + _mm_storeu_ps(d8 + 4*iy, scales128); + auto bsums4 = _mm_cvtepi32_ps(_mm_cvtepi16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][i4].d+4)))); + bsums4 = _mm_mul_ps(bsums4, scales128); + auto bsums256 = _mm256_set_m128(bsums4, bsums4); + auto bsums = _mm512_insertf32x8(_mm512_castps256_ps512(bsums256), bsums256, 1); + acc[iy] = _mm512_fmadd_ps(mx[0], _mm512_shuffle_ps(bsums, bsums, 0x00), acc[iy]); + acc[iy] = _mm512_fmadd_ps(mx[1], _mm512_shuffle_ps(bsums, bsums, 0x55), acc[iy]); + acc[iy] = _mm512_fmadd_ps(mx[2], _mm512_shuffle_ps(bsums, bsums, 0xaa), acc[iy]); + acc[iy] = _mm512_fmadd_ps(mx[3], _mm512_shuffle_ps(bsums, bsums, 0xff), acc[iy]); + } + } + for (int ib32 = 0; ib32 < 4; ++ib32) { + auto scales_l = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8l[4*i4+ib32].d)); + auto scales_h = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)q8h[4*i4+ib32].d)); + auto scales = _mm512_insertf32x8(_mm512_castps256_ps512(scales_l), scales_h, 1); + for (int j = 0; j < 8; ++j) { + qx[j] = _mm512_inserti32x8(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)q8l[4*i4+ib32].qs+j)), + _mm256_loadu_si256((const __m256i *)q8h[4*i4+ib32].qs+j), 1); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = qx_r8_q8_dot_product(qx, q8.y[iy][i4].qs+32*ib32); + auto dy = _mm512_set1_ps(d8[4*iy+ib32]); + acc[iy] = _mm512_fmadd_ps(_mm512_mul_ps(scales, dy), _mm512_cvtepi32_ps(sumi), acc[iy]); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = _mm512_setzero_ps(); + } + } + } +} +#else +template +static void mul_mat_q8_1_r8_q8_2(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + int nb = n / QK8_0; + __m256 acc[nrc_y] = {}; + float d8[4*nrc_y]; + __m256i qx[4]; + auto dot = [&qx] (const int8_t * qy) { + auto y128 = _mm_loadu_si128((const __m128i*)qy); + auto y = MM256_SET_M128I(y128, y128); +#ifdef HAVE_VNNI256 + auto sumi = _mm256_setzero_si256(); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[0], _mm256_shuffle_epi32(y, 0x00)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[1], _mm256_shuffle_epi32(y, 0x55)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[2], _mm256_shuffle_epi32(y, 0xaa)); + sumi = ggml_mm256_dpbusd_epi32(sumi, qx[3], _mm256_shuffle_epi32(y, 0xff)); + return sumi; +#else + auto sumi1 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[0], _mm256_shuffle_epi32(y, 0x00)), + _mm256_maddubs_epi16(qx[1], _mm256_shuffle_epi32(y, 0x55))); + auto sumi2 = _mm256_add_epi16(_mm256_maddubs_epi16(qx[2], _mm256_shuffle_epi32(y, 0xaa)), + _mm256_maddubs_epi16(qx[3], _mm256_shuffle_epi32(y, 0xff))); + return _mm256_add_epi32(_mm256_madd_epi16(_mm256_set1_epi16(1), sumi1), _mm256_madd_epi16(_mm256_set1_epi16(1), sumi2)); +#endif + }; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q8_1_r8 * iq8 = (const block_q8_1_r8 *)((const char *)vx + ix*bx); + for (int i4 = 0; i4 < nb/4; ++i4) { + { + __m256 mx[4]; + for (int ib32 = 0; ib32 < 4; ++ib32) mx[ib32] = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[4*i4+ib32].d+1)); + for (int iy = 0; iy < nrc_y; ++iy) { + auto scales = _mm_castsi128_ps(_mm_slli_epi32(_mm_cvtepu16_epi32(_mm_loadl_epi64((const __m128i *)q8.y[iy][i4].d)), 16)); + _mm_storeu_ps(d8 + 4*iy + 0, scales); + auto bsums4 = _mm_cvtepi32_ps(_mm_cvtepi16_epi32(_mm_loadl_epi64((const __m128i *)(q8.y[iy][i4].d+4)))); + bsums4 = _mm_mul_ps(bsums4, scales); + auto bsums = _mm256_set_m128(bsums4, bsums4); + acc[iy] = _mm256_fmadd_ps(mx[0], _mm256_shuffle_ps(bsums, bsums, 0x00), acc[iy]); + acc[iy] = _mm256_fmadd_ps(mx[1], _mm256_shuffle_ps(bsums, bsums, 0x55), acc[iy]); + acc[iy] = _mm256_fmadd_ps(mx[2], _mm256_shuffle_ps(bsums, bsums, 0xaa), acc[iy]); + acc[iy] = _mm256_fmadd_ps(mx[3], _mm256_shuffle_ps(bsums, bsums, 0xff), acc[iy]); + } + } + for (int ib32 = 0; ib32 < 4; ++ib32) { + auto scales = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)iq8[4*i4+ib32].d)); + for (int j = 0; j < 4; ++j) { + qx[j] = _mm256_loadu_si256((const __m256i *)iq8[4*i4+ib32].qs+j); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(q8.y[iy][i4].qs+32*ib32); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[4*iy+ib32])); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + } + for (int j = 0; j < 4; ++j) { + qx[j] = _mm256_loadu_si256((const __m256i *)iq8[4*i4+ib32].qs+4+j); + } + for (int iy = 0; iy < nrc_y; ++iy) { + auto sumi = dot(q8.y[iy][i4].qs+32*ib32+16); + auto d4d8 = _mm256_mul_ps(scales, _mm256_set1_ps(d8[4*iy+ib32])); + acc[iy] = _mm256_fmadd_ps(d4d8, _mm256_cvtepi32_ps(sumi), acc[iy]); + } + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, acc[iy]); + acc[iy] = _mm256_setzero_ps(); + } + } +} +#endif + +void iqk_convert_q80_q80_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + static_assert(QK4_0 == QK8_0); + GGML_ASSERT(n%QK4_0 == 0); + GGML_ASSERT(nrc_x%8 == 0); + + const int nb = n/QK4_0; + + block_q8_0_r8 * y = (block_q8_0_r8 *)vy; + + const block_q8_0 * x8[8]; + + uint32_t block[8]; + + for (int ix = 0; ix < nrc_x; ix += 8) { + + for (int k = 0; k < 8; ++k) x8[k] = (const block_q8_0 *)((const char *)vx + (ix + k)*bx); + + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + y[i].d[k] = x8[k][i].d; + _mm256_storeu_si256((__m256i *)block, _mm256_loadu_si256((const __m256i *)x8[k][i].qs)); + auto qs = (uint32_t *)y[i].qs; + for (int l = 0; l < 4; ++l) { + qs[8*l + k + 0] = block[l + 0]; + qs[8*l + k + 32] = block[l + 4]; + } + } + } + y += nb; + } +} + +template +void iqk_convert_qX_q80_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK4_0 == 0); + GGML_ASSERT(nrc_x%8 == 0); + + const int nb = n/QK8_0; + + block_q8_0_r8 * y = (block_q8_0_r8 *)vy; + + const Block * x8[8]; + + uint32_t block[8]; + + Dequantizer deq; + + for (int ix = 0; ix < nrc_x; ix += 8) { + + for (int k = 0; k < 8; ++k) x8[k] = (const Block *)((const char *)vx + (ix + k)*bx); + + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + if constexpr (std::is_same_v) { + y[i].d[k] = GGML_FP32_TO_FP16(GGML_E8M0_TO_FP32_HALF(x8[k][i].e)); + } else { + y[i].d[k] = x8[k][i].d; + } + _mm256_storeu_si256((__m256i *)block, deq.dequant(x8[k] + i)); + auto qs = (uint32_t *)y[i].qs; + for (int l = 0; l < 4; ++l) { + qs[8*l + k + 0] = block[l + 0]; + qs[8*l + k + 32] = block[l + 4]; + } + } + } + y += nb; + } +} + +template +void iqk_convert_qX_1_q8_1_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK8_0 == 0); + GGML_ASSERT(nrc_x%8 == 0); + + int nb = n/QK8_0; + + const Block * x8[8]; + + block_q8_1_r8 * y = (block_q8_1_r8 *)vy; + + uint32_t block[8]; + + Dequantizer deq; + + for (int ix = 0; ix < nrc_x; ix += 8) { + for (int k = 0; k < 8; ++k) x8[k] = (const Block *)((const char *)vx + (ix + k)*bx); + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + y[i].d[k+0] = x8[k][i].d; + y[i].d[k+8] = x8[k][i].m; + _mm256_storeu_si256((__m256i *)block, deq.dequant(x8[k]+i)); + auto qs = (uint32_t *)y[i].qs; + for (int l = 0; l < 4; ++l) { + qs[8*l + k + 0] = block[l + 0]; + qs[8*l + k + 32] = block[l + 4]; + } + } + } + y += nb; + } +} + +template void set_functions(std::array& funcs) { + if constexpr (std::is_same_v || std::is_same_v) { + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_0_q8_0_T, Dequantizer, funcs) + } + else if constexpr (std::is_same_v) { + IQK_SET_MUL_MAT_FUNCTIONS_T2(mul_mat_qX_0_q8_0_T, Dequantizer, block_q8_2, funcs) + } + else if constexpr (std::is_same_v || std::is_same_v) { + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_1_q8_2_T, Dequantizer, funcs) + } + else if constexpr (std::is_same_v) { + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_1_q8_2_T, Dequantizer, funcs) + } + else if constexpr (std::is_same_v) { + IQK_SET_MUL_MAT_FUNCTIONS_T2(mul_mat_qX_0_q8_0_T, Dequantizer, block_q8_2, funcs) + } + else if constexpr (std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v) { + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_1_q8_2_T, Dequantizer, funcs) + } +} + +} // namespace + +bool iqk_convert_legacy_quants_q8_r8(int type, int n, const void * vx, size_t bx, void * vy, int nrc_x) { + switch (type) { + case GGML_TYPE_Q4_0 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q4_1 : iqk_convert_qX_1_q8_1_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q5_0 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q5_1 : iqk_convert_qX_1_q8_1_r8>(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q6_0 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_IQ4_NL: iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q8_0 : iqk_convert_q80_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_MXFP4 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + default: return false; + } + return true; +} + +bool iqk_set_kernels_legacy_quants(int ne00, int typeA, int typeB, std::array& kernels, mul_mat_t& func16) { + + if (ne00%QK8_0 != 0) return false; + + auto expected_typeB = GGML_TYPE_Q8_2_X4; + + func16 = nullptr; + + switch (typeA) { + case GGML_TYPE_Q4_0: + set_functions(kernels); + break; + case GGML_TYPE_Q4_1: + set_functions(kernels); + break; + case GGML_TYPE_Q5_0: + set_functions(kernels); + break; + case GGML_TYPE_Q5_1: + set_functions(kernels); + break; + case GGML_TYPE_Q6_0: + set_functions(kernels); + break; + case GGML_TYPE_Q8_0: +#ifdef HAVE_FANCY_SIMD + set_functions(kernels); +#else + set_functions(kernels); +#endif + break; + case GGML_TYPE_IQ4_NL: +#ifdef HAVE_FANCY_SIMD + set_functions(kernels); +#else + set_functions(kernels); +#endif + break; + case GGML_TYPE_MXFP4: + set_functions(kernels); + break; + case GGML_TYPE_Q4_0_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q4_0_r8_q8_2, kernels) +#ifdef HAVE_FANCY_SIMD + func16 = mul_mat_q4_0_r8_q8_2<16>; +#endif + break; + case GGML_TYPE_Q5_0_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q5_0_r4_q8_2, kernels) + break; + case GGML_TYPE_Q6_0_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q6_0_r4_q8_2, kernels) + break; + case GGML_TYPE_Q8_0_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_0_r8_q8_2, kernels) + break; + case GGML_TYPE_IQ4_NL_R4: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_iq4_nl_r4_q8_2, kernels) + break; + case GGML_TYPE_Q8_1: // Note: we are misusing the Q8_1 type for Q8_1_R8 + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_1_r8_q8_2, kernels) + break; + default: + return false; + } + + return ggml_type(typeB) == expected_typeB; +} + +#else +// ---------------------------- __aarch64__ ---------------------------------------------- + +namespace { + +template +inline float16x4_t load_scales_q0(const Block * x, ggml_half * aux) { + for (int k = 0; k < 4; ++k) aux[k] = x[k].d; + return vld1_f16((const float16_t *)aux); +} + +template +inline float16x8_t load_scales_q1(const Block * x, ggml_half * aux) { + if constexpr (std::is_same_v) { + for (int k = 0; k < 4; ++k) { aux[k] = x[k].d; aux[k+4] = x[k].s; } + } else { + for (int k = 0; k < 4; ++k) { aux[k] = x[k].d; aux[k+4] = x[k].m; } + } + return vld1q_f16((const float16_t *)aux); +} + +struct Q4LegacyBits { + template + inline void prepare(const Block * x) { + for (int i = 0; i < 4; ++i) { + auto q4bits = vld1q_u8(x[i].qs); + b[2*i+0] = vreinterpretq_s8_u8(vandq_u8(q4bits, m4b)); + b[2*i+1] = vreinterpretq_s8_u8(vshrq_n_u8(q4bits, 4)); + } + } + inline void prepare1(const uint8_t * qs, int8x16_t * q) const { + auto q4bits = vld1q_u8(qs); + q[0] = vreinterpretq_s8_u8(vandq_u8(q4bits, m4b)); + q[1] = vreinterpretq_s8_u8(vshrq_n_u8(q4bits, 4)); + } + inline void prepare1(const uint8_t * qs) { + prepare1(qs, b); + } + const uint8x16_t m4b = vdupq_n_u8(0xf); + int8x16_t b[8]; +}; + +// One would think this commented out version would do better than the one below +// because it offers more opportunities to execute instructions in parallel. +// Instead, it runs significantly slower. Why? If the compiler is running out of vector registers +// cannot it just do the sequential version below on its own? +//inline int32x4_t sum_4_blocks(const int8x16_t * b, const int8_t * qs) { +// const auto q8b_1 = vld1q_s8_x2(qs + 0); +// auto p12 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b[0], q8b_1.val[0]), b[1], q8b_1.val[1]); +// const auto q8b_2 = vld1q_s8_x2(qs + 32); +// auto p34 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b[2], q8b_2.val[0]), b[3], q8b_2.val[1]); +// auto p1234 = vpaddq_s32(p12, p34); +// const auto q8b_3 = vld1q_s8_x2(qs + 64); +// auto p56 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b[4], q8b_3.val[0]), b[5], q8b_3.val[1]); +// const auto q8b_4 = vld1q_s8_x2(qs + 96); +// auto p78 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b[6], q8b_4.val[0]), b[7], q8b_4.val[1]); +// return vpaddq_s32(p1234, vpaddq_s32(p56, p78)); +//} + +inline int32x4_t sum_4_blocks(const int8x16_t * b, const int8_t * qs) { + auto q8b = vld1q_s8_x2(qs + 0); + auto p12 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b[0], q8b.val[0]), b[1], q8b.val[1]); + q8b = vld1q_s8_x2(qs + 32); + auto p34 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b[2], q8b.val[0]), b[3], q8b.val[1]); + auto p1234 = vpaddq_s32(p12, p34); + q8b = vld1q_s8_x2(qs + 64); + auto p56 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b[4], q8b.val[0]), b[5], q8b.val[1]); + q8b = vld1q_s8_x2(qs + 96); + auto p78 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b[6], q8b.val[0]), b[7], q8b.val[1]); + return vpaddq_s32(p1234, vpaddq_s32(p56, p78)); +} + +inline int32x4x2_t sum_4_blocks(const int8x16_t * b1, const int8x16_t * b2, const int8_t * qs) { + auto q8b = vld1q_s8_x2(qs + 0); + auto p12_1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b1[0], q8b.val[0]), b1[1], q8b.val[1]); + auto p12_2 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b2[0], q8b.val[0]), b2[1], q8b.val[1]); + q8b = vld1q_s8_x2(qs + 32); + auto p34_1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b1[2], q8b.val[0]), b1[3], q8b.val[1]); + auto p34_2 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b2[2], q8b.val[0]), b2[3], q8b.val[1]); + auto p1234_1 = vpaddq_s32(p12_1, p34_1); + auto p1234_2 = vpaddq_s32(p12_2, p34_2); + q8b = vld1q_s8_x2(qs + 64); + auto p56_1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b1[4], q8b.val[0]), b1[5], q8b.val[1]); + auto p56_2 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b2[4], q8b.val[0]), b2[5], q8b.val[1]); + q8b = vld1q_s8_x2(qs + 96); + auto p78_1 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b1[6], q8b.val[0]), b1[7], q8b.val[1]); + auto p78_2 = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), b2[6], q8b.val[0]), b2[7], q8b.val[1]); + auto p5678_1 = vpaddq_s32(p56_1, p78_1); + auto p5678_2 = vpaddq_s32(p56_2, p78_2); + return { vpaddq_s32(p1234_1, p5678_1), vpaddq_s32(p1234_2, p5678_2)}; +} + +template struct Q80 { + + constexpr static int nrc_y = nrc; + + Q80(const DataInfo& info) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const block_q8_0 *)info.src1_row(iy); + } + + inline const int8_t * quant_data(int iy, int i) const { + const block_q8_0_x4 * y4 = (const block_q8_0_x4 *)y[iy] + i; + return y4->qs; + } + + inline float16x4_t load_scales(int iy, int i) const { + const block_q8_0_x4 * y4 = (const block_q8_0_x4 *)y[iy] + i; + return vld1_f16((const float16_t *)y4->d); + } + + template + inline void process_scales(int i, Dequantizer& deq, float16x4_t * sc16, float32x4_t * /*acc*/) const { + auto qx_scales = deq.new_block(i); + for (int iy = 0; iy < nrc; ++iy) { + auto q8_scales = load_scales(iy, i); + sc16[iy] = vmul_f16(qx_scales, q8_scales); + } + } + + template + inline void process_scales(int i, Dequantizer& deq1, Dequantizer& deq2, float16x4_t * sc16, float32x4_t * /*acc*/) const { + auto qx_scales_1 = deq1.new_block(i); + auto qx_scales_2 = deq2.new_block(i); + for (int iy = 0; iy < nrc; ++iy) { + auto q8_scales = load_scales(iy, i); + sc16[iy ] = vmul_f16(qx_scales_1, q8_scales); + sc16[iy+nrc_y] = vmul_f16(qx_scales_2, q8_scales); + } + } + + template + inline void process_1_block(int i, Dequantizer& deq, float32x4_t * acc) const { + deq.prepare1(i); + float d = deq.block_scale(i); + for (int iy = 0; iy < nrc; ++iy) { + auto q8b = vld1q_s8_x2(y[iy][i].qs); + auto p = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b[0], q8b.val[0]), deq.bits.b[1], q8b.val[1]); + acc[iy] = vmlaq_f32(acc[iy], vdupq_n_f32(d*GGML_FP16_TO_FP32(y[iy][i].d)), vcvtq_f32_s32(p)); + } + } + + const block_q8_0 * y[nrc_y]; +}; + +template struct Q81 { + + constexpr static int nrc_y = nrc; + + Q81(const DataInfo& info) { + for (int iy = 0; iy < nrc_y; ++iy) y[iy] = (const block_q8_1 *)info.src1_row(iy); + } + + inline const int8_t * quant_data(int iy, int i) const { + const block_q8_1_x4 * y4 = (const block_q8_1_x4 *)y[iy] + i; + return y4->qs; + } + + inline float16x8_t load_scales(int iy, int i) const { + const block_q8_1_x4 * y4 = (const block_q8_1_x4 *)y[iy] + i; + return vld1q_f16((const float16_t *)y4->d); + } + + template + inline void process_scales(int i, Dequantizer& deq, float16x4_t * sc16, float32x4_t * acc) const { + auto qx_scales = deq.new_block(i); + for (int iy = 0; iy < nrc; ++iy) { + auto q8_scales = load_scales(iy, i); + auto m = vmul_f16(vget_high_f16(qx_scales), vget_high_f16(q8_scales)); + acc[iy] = vaddq_f32(acc[iy], vcvt_f32_f16(m)); + sc16[iy] = vmul_f16(vget_low_f16(qx_scales), vget_low_f16(q8_scales)); + } + } + + template + inline void process_scales(int i, Dequantizer& deq1, Dequantizer& deq2, float16x4_t * sc16, float32x4_t * acc) const { + auto qx_scales_1 = deq1.new_block(i); + auto qx_scales_2 = deq2.new_block(i); + for (int iy = 0; iy < nrc; ++iy) { + auto q8_scales = load_scales(iy, i); + auto q8_scales_l = vget_low_f16(q8_scales); + auto q8_scales_h = vget_high_f16(q8_scales); + auto m1 = vmul_f16(vget_high_f16(qx_scales_1), q8_scales_h); + auto m2 = vmul_f16(vget_high_f16(qx_scales_2), q8_scales_h); + acc[iy ] = vaddq_f32(acc[iy ], vcvt_f32_f16(m1)); + acc[iy+nrc_y ] = vaddq_f32(acc[iy+nrc_y], vcvt_f32_f16(m2)); + sc16[iy ] = vmul_f16(vget_low_f16(qx_scales_1), q8_scales_l); + sc16[iy+nrc_y] = vmul_f16(vget_low_f16(qx_scales_2), q8_scales_l); + } + } + + template + inline void process_1_block(int i, Dequantizer& deq, float32x4_t * acc) const { + deq.prepare1(i); + float d = GGML_FP16_TO_FP32(deq.x[i].d), m = 0.25f*GGML_FP16_TO_FP32(deq.x[i].m); + for (int iy = 0; iy < nrc; ++iy) { + auto q8b = vld1q_s8_x2(y[iy][i].qs); + auto p = ggml_vdotq_s32(ggml_vdotq_s32(vdupq_n_s32(0), deq.bits.b[0], q8b.val[0]), deq.bits.b[1], q8b.val[1]); + acc[iy] = vmlaq_f32(acc[iy], vdupq_n_f32(d*GGML_FP16_TO_FP32(y[iy][i].d)), vcvtq_f32_s32(p)); + acc[iy] = vaddq_f32(acc[iy], vdupq_n_f32(m*GGML_FP16_TO_FP32(y[iy][i].s))); + } + } + + const block_q8_1 * y[nrc_y]; +}; + +template +struct BaseLegacyDequantizer { + + BaseLegacyDequantizer(const void * vx, size_t bx) : vx(vx), x(nullptr), bx(bx) {} + + inline void new_row(int ix) { x = (const block_q *)((const char *)vx + bx*ix); } + + Q4LegacyBits bits; + + const void * vx; + const block_q * x; + size_t bx; +}; + +struct DequantizerQ40 final : public BaseLegacyDequantizer { + + DequantizerQ40(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i, int8x16_t * q) const { + bits.prepare1(x[i].qs, q); + q[0] = vaddq_s8(q[0], m8); + q[1] = vaddq_s8(q[1], m8); + } + inline void prepare1(int i) { + prepare1(i, bits.b); + } + + inline float16x4_t new_block(int i) { + ggml_half aux[4]; + for (int k = 0; k < 4; ++k) { + aux[k] = x[4*i+k].d; + prepare1(4*i+k, bits.b + 2*k); + } + return vld1_f16((const float16_t *)aux); + } + + inline float block_scale(int i) const { return GGML_FP16_TO_FP32(x[i].d); } + + const int8x16_t m8 = vdupq_n_s8(-8); + //ggml_half aux[4]; +}; + +struct DequantizerQ60 final : public BaseLegacyDequantizer { + + DequantizerQ60(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i, int8x16_t * q) const { + bits.prepare1(x[i].qs, q); + auto qh8 = vld1_u8(x[i].qh); + auto qh = vcombine_u8(vshl_n_u8(qh8, 4), qh8); + q[0] = vaddq_s8(vorrq_u8(q[0], vandq_u8(qh, hmask)), m32); + q[1] = vaddq_s8(vorrq_u8(q[1], vandq_u8(vshrq_n_u8(qh, 2), hmask)), m32); + } + inline void prepare1(int i) { + prepare1(i, bits.b); + } + + inline float16x4_t new_block(int i) { + ggml_half aux[4]; + for (int k = 0; k < 4; ++k) { + aux[k] = x[4*i+k].d; + prepare1(4*i+k, bits.b + 2*k); + } + return vld1_f16((const float16_t *)aux); + } + inline float block_scale(int i) const { return GGML_FP16_TO_FP32(x[i].d); } + + const int8x16_t m32 = vdupq_n_s8(-32); + const uint8x16_t hmask = vdupq_n_u8(0x30); +}; + +struct DequantizerIQ4NL final : public BaseLegacyDequantizer { + + DequantizerIQ4NL(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i, int8x16_t * q) const { + bits.prepare1(x[i].qs, q); + q[0] = vqtbl1q_s8(values, q[0]); + q[1] = vqtbl1q_s8(values, q[1]); + } + inline void prepare1(int i) { + prepare1(i, bits.b); + } + + inline float16x4_t new_block(int i) { + ggml_half aux[4]; + for (int k = 0; k < 4; ++k) { + aux[k] = x[4*i+k].d; + prepare1(4*i+k, bits.b + 2*k); + } + return vld1_f16((const float16_t *)aux); + } + static int8x16_t load_values() { + static const int8_t iq4nl_values[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; + return vld1q_s8(iq4nl_values); + } + inline float block_scale(int i) const { return GGML_FP16_TO_FP32(x[i].d); } + + const int8x16_t values = load_values(); +}; + +struct DequantizerMXFP4 final : public BaseLegacyDequantizer { + + DequantizerMXFP4(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i, int8x16_t * q) const { + bits.prepare1(x[i].qs, q); + q[0] = vqtbl1q_s8(values, q[0]); + q[1] = vqtbl1q_s8(values, q[1]); + } + inline void prepare1(int i) { + prepare1(i, bits.b); + } + + inline float16x4_t new_block(int i) { + float aux[4]; + for (int k = 0; k < 4; ++k) { + aux[k] = GGML_E8M0_TO_FP32_HALF(x[4*i+k].e); + prepare1(4*i+k, bits.b + 2*k); + } + return vcvt_f16_f32(vld1q_f32(aux)); + } + static int8x16_t load_values() { + return vld1q_s8(kvalues_mxfp4); + } + inline float block_scale(int i) const { return GGML_E8M0_TO_FP32_HALF(x[i].e); } + + const int8x16_t values = load_values(); +}; + +struct DequantizerQ41 : public BaseLegacyDequantizer { + + DequantizerQ41(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i) { + bits.prepare1(x[i].qs); + } + + inline float16x8_t new_block(int i) { + uint32_t aux32[4]; + const uint32_t * s32 = (const uint32_t *)&x[4*i].d; + for (int k = 0; k < 4; ++k) { + aux32[k] = *s32; s32 += sizeof(block_q4_1)/4; + bits.prepare1(x[4*i+k].qs, bits.b + 2*k); + } + return vreinterpretq_f16_u8(vqtbl1q_u8(vld1q_u8((const uint8_t *)aux32), vreinterpretq_u8_u64(shuffle))); + } + // Leaving this commented out attempt to be reminded that I already tried this. + // It has basically the same performance as the version above. + //inline float16x8_t new_block(int i) { + // uint32x4_t scales = {}; + // const block_q4_1 * xi = x + 4*i; + // const uint32_t * s32 = (const uint32_t *)&xi->d; + // scales = vsetq_lane_u32(*s32, scales, 0); s32 += sizeof(block_q4_1)/4; + // bits.prepare1(xi[0].qs, bits.b + 0); + // scales = vsetq_lane_u32(*s32, scales, 1); s32 += sizeof(block_q4_1)/4; + // bits.prepare1(xi[1].qs, bits.b + 2); + // scales = vsetq_lane_u32(*s32, scales, 2); s32 += sizeof(block_q4_1)/4; + // bits.prepare1(xi[2].qs, bits.b + 4); + // scales = vsetq_lane_u32(*s32, scales, 3); + // bits.prepare1(xi[3].qs, bits.b + 6); + // return vreinterpretq_f16_u8(vqtbl1q_u8(vreinterpretq_u8_u32(scales), vreinterpretq_u8_u64(shuffle))); + //} + + const uint64x2_t shuffle = {0x0d0c090805040100, 0x0f0e0b0a07060302}; +}; + +struct HighBit5Legacy { + inline uint8x16_t to_bytes(const uint8_t * qh) const { + uint8x16_t h = vqtbl1q_u8(vreinterpretq_u8_u16(vdupq_n_u16(*(const uint16_t *)qh)), shuffle); + return vceqq_u8(vandq_u8(h, vreinterpretq_u8_u64(mask)), vreinterpretq_u8_u64(mask)); + } + inline uint8x16_t to_negated_bytes(const uint8_t * qh) const { + uint8x16_t h = vqtbl1q_u8(vreinterpretq_u8_u16(vdupq_n_u16(*(const uint16_t *)qh)), shuffle); + return vceqq_u8(vandq_u8(h, vreinterpretq_u8_u64(mask)), vdupq_n_u8(0)); + } + const uint64x2_t mask = vdupq_n_u64(0x8040201008040201); + const uint8x16_t shuffle = vcombine_u8(vdup_n_u8(0), vdup_n_u8(1)); +}; + +struct DequantizerQ50 final : public BaseLegacyDequantizer { + + DequantizerQ50(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i, int8x16_t * q) const { + bits.prepare1(x[i].qs, q); + auto qh = x[i].qh; + q[0] = vreinterpretq_s8_u8(vorrq_u8(vreinterpretq_u8_s8(q[0]), vandq_u8(mh, hbits.to_negated_bytes(qh+0)))); + q[1] = vreinterpretq_s8_u8(vorrq_u8(vreinterpretq_u8_s8(q[1]), vandq_u8(mh, hbits.to_negated_bytes(qh+2)))); + } + inline void prepare1(int i) { + prepare1(i, bits.b); + } + + inline float16x4_t new_block(int i) { + ggml_half aux[4]; + for (int k = 0; k < 4; ++k) { + aux[k] = x[4*i+k].d; + prepare1(4*i+k, bits.b + 2*k); + } + return vld1_f16((const float16_t *)aux); + } + inline float block_scale(int i) const { return GGML_FP16_TO_FP32(x[i].d); } + + HighBit5Legacy hbits; + + const uint8x16_t mh = vdupq_n_u8(0xf0); + +}; + +struct DequantizerQ80 final : public BaseLegacyDequantizer { + + DequantizerQ80(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i) { + bits.b[0] = vld1q_s8(x[i].qs); + bits.b[1] = vld1q_s8(x[i].qs+16); + } + + inline float16x4_t new_block(int i) { + ggml_half aux[4]; + for (int k = 0; k < 4; ++k) { + aux[k] = x[4*i+k].d; + bits.b[2*k+0] = vld1q_s8(x[4*i+k].qs); + bits.b[2*k+1] = vld1q_s8(x[4*i+k].qs+16); + } + return vld1_f16((const float16_t *)aux); + } + inline float block_scale(int i) const { return GGML_FP16_TO_FP32(x[i].d); } + +}; + +// TODO: handle case where row size is not a multiple of 128 +struct DequantizerQ80_x4 final : public BaseLegacyDequantizer { + + DequantizerQ80_x4(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i) { + bits.b[0] = vld1q_s8(x[i].qs); + bits.b[1] = vld1q_s8(x[i].qs+16); + } + + inline float16x4_t new_block(int i) { + auto scale = vld1_f16((const float16_t *)x[i].d); + for (int k = 0; k < 4; ++k) { + bits.b[2*k+0] = vld1q_s8(x[i].qs+32*k); + bits.b[2*k+1] = vld1q_s8(x[i].qs+32*k+16); + } + return scale; + } + +}; + +struct DequantizerQ51 final : public BaseLegacyDequantizer { + + DequantizerQ51(const void * vx, size_t bx) : BaseLegacyDequantizer(vx, bx) {} + + inline void prepare1(int i, int8x16_t * q) const { + bits.prepare1(x[i].qs, q); + auto qh = x[i].qh; + q[0] = vreinterpretq_s8_u8(vorrq_u8(vreinterpretq_u8_s8(q[0]), vandq_u8(mh, hbits.to_bytes(qh+0)))); + q[1] = vreinterpretq_s8_u8(vorrq_u8(vreinterpretq_u8_s8(q[1]), vandq_u8(mh, hbits.to_bytes(qh+2)))); + } + inline void prepare1(int i) { + bits.prepare1(x[i].qs, bits.b); + } + + inline float16x8_t new_block(int i) { + uint32_t aux32[4]; + const uint32_t * s32 = (const uint32_t *)&x[4*i].d; + for (int k = 0; k < 4; ++k) { + aux32[k] = *s32; s32 += sizeof(block_q5_1)/4; + prepare1(4*i+k, bits.b + 2*k); + } + return vreinterpretq_f16_u8(vqtbl1q_u8(vld1q_u8((const uint8_t *)aux32), vreinterpretq_u8_u64(shuffle))); + } + + HighBit5Legacy hbits; + + const uint8x16_t mh = vdupq_n_u8(0x10); + const uint64x2_t shuffle = {0x0d0c090805040100, 0x0f0e0b0a07060302}; + +}; + +template +inline void sum_4(int i, Dequantizer& deq, const Q8& q8, const float16x4_t * sc16, float32x4_t * acc) { + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + auto pall = sum_4_blocks(deq.bits.b, q8.quant_data(iy, i)); + auto scale = vcvt_f32_f16(sc16[iy]); + acc[iy] = vmlaq_f32(acc[iy], scale, vcvtq_f32_s32(pall)); + } +} + +template +inline void sum_4(int i, Dequantizer& deq1, Dequantizer& deq2, const Q8& q8, const float16x4_t * sc16, float32x4_t * acc) { + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + auto pall = sum_4_blocks(deq1.bits.b, deq2.bits.b, q8.quant_data(iy, i)); + auto scale1 = vcvt_f32_f16(sc16[iy]); + auto scale2 = vcvt_f32_f16(sc16[iy+Q8::nrc_y]); + acc[iy] = vmlaq_f32(acc[iy], scale1, vcvtq_f32_s32(pall.val[0])); + acc[iy+Q8::nrc_y] = vmlaq_f32(acc[iy+Q8::nrc_y], scale2, vcvtq_f32_s32(pall.val[1])); + } +} + +template +inline void mul_mat_qX_Y_q8_Y(int n, Dequantizer& deq, Q8& q8, const DataInfo& info, int nrc_x) { + const int nb = n / QK4_1; + + float16x4_t sc16[Q8::nrc_y]; + + for (int ix = 0; ix < nrc_x; ++ix) { + + deq.new_row(ix); + + float32x4_t acc[Q8::nrc_y]; + for (int iy = 0; iy < Q8::nrc_y; ++iy) acc[iy] = vdupq_n_f32(0.f); + + for (int i = 0; i < nb/4; ++i) { + q8.process_scales(i, deq, sc16, acc); + sum_4(i, deq, q8, sc16, acc); + } + for (int i = 4*(nb/4); i < nb; ++i) { + q8.process_1_block(i, deq, acc); + } + + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + info.store(ix, iy, vaddvq_f32(acc[iy])); + } + } +} + +template +inline void mul_mat_qX_Y_q8_Y_IK(int n, Dequantizer& deq1, Dequantizer& deq2, Q8& q8, const DataInfo& info, int nrc_x) { + const int nb = n / QK4_1; + + float16x4_t sc16[2*Q8::nrc_y]; + float32x4_t acc[2*Q8::nrc_y]; + + for (int ix = 0; ix < nrc_x; ix += 2) { + + deq1.new_row(ix+0); + deq2.new_row(ix+1); + + for (int iy = 0; iy < 2*Q8::nrc_y; ++iy) acc[iy] = vdupq_n_f32(0.f); + + for (int i = 0; i < nb/4; ++i) { + q8.process_scales(i, deq1, deq2, sc16, acc); + sum_4(i, deq1, deq2, q8, sc16, acc); + } + //for (int i = 4*(nb/4); i < nb; ++i) { + // q8.process_1_block(i, deq, acc); + //} + + for (int iy = 0; iy < Q8::nrc_y; ++iy) { + info.store(ix+0, iy, vaddvq_f32(acc[iy])); + info.store(ix+1, iy, vaddvq_f32(acc[iy+Q8::nrc_y])); + } + } +} + +template +inline void mul_mat_qX_Y_q8_Y_1(int n, Dequantizer& deq1, Dequantizer& deq2, Q8& q8, const DataInfo& info, int nrc_x) { + const int nb = n / QK4_1; + + float16x4_t sc16[2]; + + for (int ix = 0; ix < nrc_x; ++ix) { + + deq1.new_row(ix); + deq2.new_row(ix); + + float32x4_t acc[2] = { vdupq_n_f32(0.f), vdupq_n_f32(0.f) }; + + for (int i = 0; i < nb/8; ++i) { + q8.process_scales(2*i+0, deq1, sc16+0, acc+0); + q8.process_scales(2*i+1, deq2, sc16+1, acc+1); + sum_4(2*i+0, deq1, q8, sc16+0, acc+0); + sum_4(2*i+1, deq2, q8, sc16+1, acc+1); + } + for (int i = 2*(nb/8); i < nb/4; ++i) { + q8.process_scales(i, deq1, sc16, acc); + sum_4(i, deq1, q8, sc16, acc); + } + //for (int i = 4*(nb/4); i < nb; ++i) { + // q8.process_1_block(i, deq1, acc); + //} + + info.store(ix, 0, vaddvq_f32(vaddq_f32(acc[0], acc[1]))); + } +} + +template +static void mul_mat_qX_1_q8_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + Q81 q8(info); + if constexpr (nrc_y == 1) { + Dequantizer deq1(vx, bx), deq2(vx, bx); + mul_mat_qX_Y_q8_Y_1(n, deq1, deq2, q8, info, nrc_x); + } else { + if (nrc_x%2 == 0 && n%128 == 0) { + Dequantizer deq1(vx, bx), deq2(vx, bx); + mul_mat_qX_Y_q8_Y_IK(n, deq1, deq2, q8, info, nrc_x); + } else { + Dequantizer deq(vx, bx); + mul_mat_qX_Y_q8_Y(n, deq, q8, info, nrc_x); + } + } +} + +template +static void mul_mat_qX_0_q8_0(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + Q80 q8(info); + if constexpr (nrc_y == 1) { + Dequantizer deq1(vx, bx), deq2(vx, bx); + mul_mat_qX_Y_q8_Y_1(n, deq1, deq2, q8, info, nrc_x); + } else { + if (nrc_x%2 == 0 && n%128 == 0) { + Dequantizer deq1(vx, bx), deq2(vx, bx); + mul_mat_qX_Y_q8_Y_IK(n, deq1, deq2, q8, info, nrc_x); + } else { + Dequantizer deq(vx, bx); + mul_mat_qX_Y_q8_Y(n, deq, q8, info, nrc_x); + } + } +} + +template +static void mul_mat_qX_1_q8_1_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + Dequantizer deq1(vx, bx), deq2(vx, bx); + Q81<1> q8(info); + mul_mat_qX_Y_q8_Y_1(n, deq1, deq2, q8, info, nrc_x); +} + +template +static void mul_mat_qX_0_q8_0_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + Dequantizer deq1(vx, bx), deq2(vx, bx); + Q80<1> q8(info); + mul_mat_qX_Y_q8_Y(n, deq1, deq2, q8, info, nrc_x); +} + +template +void mul_mat_qx_r4_q8_0(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%4 == 0); + Q8 q8(info); + Dequantizer deq(vx, bx); + int nb = n / QK4_NL; + int8x16_t qx[8]; + float d8[4*nrc_y]; + float32x4_t acc[nrc_y] = {}; + for (int ix = 0; ix < nrc_x; ix += 4) { + deq.new_row(ix); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + vst1q_f32(d8+4*iy, vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][ib4].d))); + } + for (int k = 0; k < 4; ++k) { + auto scales = deq.prepare(4*ib4+k, qx); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][ib4].qs+32*k); + auto sumi = interleaved_dotq(qx, y); + auto d4d8 = vmulq_f32(scales, vdupq_n_f32(d8[4*iy+k])); + acc[iy] = vfmaq_f32(acc[iy], d4d8, vcvtq_f32_s32(sumi)); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = deq.prepare(ib, qx); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_0 *)q8.y[iy]; + auto y = vld1q_s8_x2(qy[ib].qs); + auto sumi = interleaved_dotq(qx, y); + auto d4d8 = vmulq_f32(scales, vdupq_n_f32(GGML_FP16_TO_FP32(qy[ib].d))); + acc[iy] = vfmaq_f32(acc[iy], d4d8, vcvtq_f32_s32(sumi)); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix, iy, deq.result(acc[iy])); + acc[iy] = vdupq_n_f32(0.f); + } + } +} + +template +void mul_mat_qx_r8_q8_0(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + Dequantizer deq(vx, bx); + int nb = n / QK4_NL; + int8x16_t qx[16]; + float d8[4*nrc_y]; + float32x4_t acc[2*nrc_y] = {}; + for (int ix = 0; ix < nrc_x; ix += 8) { + deq.new_row(ix); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + vst1q_f32(d8+4*iy, vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][ib4].d))); + } + for (int k = 0; k < 4; ++k) { + auto scales = deq.prepare(ib4, k, qx); + for (int iy = 0; iy < nrc_y; ++iy) { + auto y = vld1q_s8_x2(q8.y[iy][ib4].qs+32*k); + auto sumi1 = interleaved_dotq(qx+0, y); + auto sumi2 = interleaved_dotq(qx+8, y); + auto dy = vdupq_n_f32(d8[4*iy+k]); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], vmulq_f32(scales.val[0], dy), vcvtq_f32_s32(sumi1)); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], vmulq_f32(scales.val[1], dy), vcvtq_f32_s32(sumi2)); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales = deq.prepare(ib, 0, qx); + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_0 *)q8.y[iy]; + auto y = vld1q_s8_x2(qy[ib].qs); + auto sumi1 = interleaved_dotq(qx+0, y); + auto sumi2 = interleaved_dotq(qx+8, y); + auto dy = vdupq_n_f32(GGML_FP16_TO_FP32(qy[ib].d)); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], vmulq_f32(scales.val[0], dy), vcvtq_f32_s32(sumi1)); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], vmulq_f32(scales.val[1], dy), vcvtq_f32_s32(sumi2)); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix+0, iy, deq.result(acc[2*iy+0])); + info.store(ix+4, iy, deq.result(acc[2*iy+1])); + acc[2*iy] = acc[2*iy+1] = vdupq_n_f32(0.f); + } + } +} + +struct IQ4_NL_R4_Dequantizer { + IQ4_NL_R4_Dequantizer(const void * vx, size_t bx) : cx((const char *)vx), bx(bx), values(vld1q_s8(iq4k_values)) {} + inline void new_row(int ix) { iq4 = (const block_iq4_nl_r4 *)(cx + ix*bx); } + inline float32x4_t prepare(int ib, int8x16_t * qx) const { + auto scales = vcvt_f32_f16(vld1_f16((const float16_t *)iq4[ib].d)); + auto bits = vld1q_u8_x4(iq4[ib].qs); + prepare_iq4_nl_quants(values, m4, bits, qx); + return scales; + } + inline float32x4_t result(float32x4_t acc) const { + return acc; + } + + const char * cx; + const size_t bx; + const block_iq4_nl_r4 * iq4; + const uint8x16_t m4 = vdupq_n_u8(0x0f); + const int8x16_t values; +}; + +struct Q4_0_R8_Dequantizer { + Q4_0_R8_Dequantizer(const void * vx, size_t bx) : cx((const char *)vx), bx(bx) {} + inline void new_row(int ix) { iq4 = (const block_iq4_nl_r8 *)(cx + ix*bx); } + inline float32x4x2_t prepare(int ib4, int k, int8x16_t * qx) const { + auto scales16 = vld1q_f16((const float16_t *)iq4[4*ib4+k].d); + float32x4x2_t scales = { vcvt_f32_f16(vget_low_f16(scales16)), vcvt_f32_f16(vget_high_f16(scales16)) }; + for (int j = 0; j < 4; ++j) { + auto bits = vld1q_u8_x2(iq4[4*ib4+k].qs + 32*j); + bits.val[0] = veorq_u8(m88, bits.val[0]); + bits.val[1] = veorq_u8(m88, bits.val[1]); + qx[2*j+0] = vshlq_n_u8(bits.val[0], 4); + qx[2*j+1] = vandq_u8(bits.val[0], m4); + qx[2*j+8] = vshlq_n_u8(bits.val[1], 4); + qx[2*j+9] = vandq_u8(bits.val[1], m4); + } + return scales; + } + inline float32x4_t result(float32x4_t acc) const { + return vmulq_f32(norm, acc); + } + + const char * cx; + const size_t bx; + const block_iq4_nl_r8 * iq4; + const uint8x16_t m4 = vdupq_n_u8(0xf0); + const uint8x16_t m88 = vdupq_n_u8(0x88); + const float32x4_t norm = vdupq_n_f32(1.f/16); +}; + +struct Q5_0_R4_Dequantizer { + Q5_0_R4_Dequantizer(const void * vx, size_t bx) : cx((const char *)vx), bx(bx) {} + inline void new_row(int ix) { iq5 = (const block_q5_0_r4 *)(cx + ix*bx); } + inline float32x4_t prepare(int ib, int8x16_t * qx) const { + auto scales = vcvt_f32_f16(vld1_f16((const float16_t *)iq5[ib].d)); + auto lbits = vld1q_u8_x4(iq5[ib].qs); + auto hbits = vld1q_u8(iq5[ib].qh); + qx[0] = vaddq_s8(vandq_u8(lbits.val[0], m4) | vandq_u8(vshlq_n_u8(hbits, 4), m5), m16); // 0...3 + qx[1] = vaddq_s8(vandq_u8(lbits.val[1], m4) | vandq_u8(vshlq_n_u8(hbits, 3), m5), m16); // 16..19 + qx[2] = vaddq_s8(vandq_u8(lbits.val[2], m4) | vandq_u8(vshlq_n_u8(hbits, 2), m5), m16); // 4...7 + qx[3] = vaddq_s8(vandq_u8(lbits.val[3], m4) | vandq_u8(vshlq_n_u8(hbits, 1), m5), m16); // 20..23 + qx[4] = vaddq_s8(vshrq_n_u8(lbits.val[0], 4)| vandq_u8(hbits, m5), m16); // 8..11 + qx[5] = vaddq_s8(vshrq_n_u8(lbits.val[1], 4)| vandq_u8(vshrq_n_u8(hbits, 1), m5), m16); // 24..27 + qx[6] = vaddq_s8(vshrq_n_u8(lbits.val[2], 4)| vandq_u8(vshrq_n_u8(hbits, 2), m5), m16); // 12..15 + qx[7] = vaddq_s8(vshrq_n_u8(lbits.val[3], 4)| vandq_u8(vshrq_n_u8(hbits, 3), m5), m16); // 28..31 + return scales; + } + inline float32x4_t result(float32x4_t acc) const { + return acc; + } + + const char * cx; + const size_t bx; + const block_q5_0_r4 * iq5; + const uint8x16_t m4 = vdupq_n_u8(0x0f); + const uint8x16_t m5 = vdupq_n_u8(0x10); + const int8x16_t m16 = vdupq_n_s8(-16); +}; + +struct Q6_0_R4_Dequantizer { + Q6_0_R4_Dequantizer(const void * vx, size_t bx) : cx((const char *)vx), bx(bx) {} + inline void new_row(int ix) { iq6 = (const block_q6_0_r4 *)(cx + ix*bx); } + inline float32x4_t prepare(int ib, int8x16_t * qx) const { + auto scales = vcvt_f32_f16(vld1_f16((const float16_t *)iq6[ib].d)); + auto lbits = vld1q_u8_x4(iq6[ib].qs); + auto hbits = vld1q_u8_x2(iq6[ib].qh); + qx[0] = vaddq_s8(vandq_u8(lbits.val[0], m4) | vandq_u8(vshlq_n_u8(hbits.val[0], 4), m6), m32); // 0...3 + qx[1] = vaddq_s8(vandq_u8(lbits.val[1], m4) | vandq_u8(vshlq_n_u8(hbits.val[1], 4), m6), m32); // 16..19 + qx[2] = vaddq_s8(vandq_u8(lbits.val[2], m4) | vandq_u8(vshlq_n_u8(hbits.val[0], 2), m6), m32); // 4...7 + qx[3] = vaddq_s8(vandq_u8(lbits.val[3], m4) | vandq_u8(vshlq_n_u8(hbits.val[1], 2), m6), m32); // 20..23 + qx[4] = vaddq_s8(vshrq_n_u8(lbits.val[0], 4)| vandq_u8(hbits.val[0], m6), m32); // 8..11 + qx[5] = vaddq_s8(vshrq_n_u8(lbits.val[1], 4)| vandq_u8(hbits.val[1], m6), m32); // 24..27 + qx[6] = vaddq_s8(vshrq_n_u8(lbits.val[2], 4)| vandq_u8(vshrq_n_u8(hbits.val[0], 2), m6), m32); // 12..15 + qx[7] = vaddq_s8(vshrq_n_u8(lbits.val[3], 4)| vandq_u8(vshrq_n_u8(hbits.val[1], 2), m6), m32); // 28..31 + return scales; + } + inline float32x4_t result(float32x4_t acc) const { + return acc; + } + + const char * cx; + const size_t bx; + const block_q6_0_r4 * iq6; + const uint8x16_t m4 = vdupq_n_u8(0x0f); + const uint8x16_t m6 = vdupq_n_u8(0x30); + const int8x16_t m32 = vdupq_n_s8(-32); +}; + +inline void qx_0_q8_0_dot(const int8x16_t * qx, const int8_t * qy, int32x4_t& sumi1, int32x4_t& sumi2) { + auto y = vld1q_s8_x2(qy); + sumi1 = sumi2 = vdupq_n_s32(0); + sumi1 = vdotq_laneq_s32(sumi1, qx[0], y.val[0], 0); + sumi2 = vdotq_laneq_s32(sumi2, qx[1], y.val[0], 0); + sumi1 = vdotq_laneq_s32(sumi1, qx[2], y.val[0], 1); + sumi2 = vdotq_laneq_s32(sumi2, qx[3], y.val[0], 1); + sumi1 = vdotq_laneq_s32(sumi1, qx[4], y.val[0], 2); + sumi2 = vdotq_laneq_s32(sumi2, qx[5], y.val[0], 2); + sumi1 = vdotq_laneq_s32(sumi1, qx[6], y.val[0], 3); + sumi2 = vdotq_laneq_s32(sumi2, qx[7], y.val[0], 3); + sumi1 = vdotq_laneq_s32(sumi1, qx[8+0], y.val[1], 0); + sumi2 = vdotq_laneq_s32(sumi2, qx[8+1], y.val[1], 0); + sumi1 = vdotq_laneq_s32(sumi1, qx[8+2], y.val[1], 1); + sumi2 = vdotq_laneq_s32(sumi2, qx[8+3], y.val[1], 1); + sumi1 = vdotq_laneq_s32(sumi1, qx[8+4], y.val[1], 2); + sumi2 = vdotq_laneq_s32(sumi2, qx[8+5], y.val[1], 2); + sumi1 = vdotq_laneq_s32(sumi1, qx[8+6], y.val[1], 3); + sumi2 = vdotq_laneq_s32(sumi2, qx[8+7], y.val[1], 3); +} + +template +void mul_mat_q8_0_r8_q8_0(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + int nb = n / QK8_0; + float32x4_t acc[2*nrc_y] = {}; + int8x16_t qx[16]; + float d8[4*nrc_y]; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q8_0_r8 * iq8 = (const block_q8_0_r8 *)((const char *)vx + ix*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + vst1q_f32(d8+4*iy, vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][ib4].d))); + } + for (int k = 0; k < 4; ++k) { + auto scales16 = vld1q_f16((const float16_t *)iq8[4*ib4+k].d); + auto scales1 = vcvt_f32_f16(vget_low_f16 (scales16)); + auto scales2 = vcvt_f32_f16(vget_high_f16(scales16)); + for (int j = 0; j < 16; ++j) qx[j] = vld1q_s8(iq8[4*ib4+k].qs + 16*j); + int32x4_t sumi1, sumi2; + for (int iy = 0; iy < nrc_y; ++iy) { + qx_0_q8_0_dot(qx, q8.y[iy][ib4].qs+32*k, sumi1, sumi2); + auto dy = vdupq_n_f32(d8[4*iy+k]); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], vmulq_f32(scales1, dy), vcvtq_f32_s32(sumi1)); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], vmulq_f32(scales2, dy), vcvtq_f32_s32(sumi2)); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales16 = vld1q_f16((const float16_t *)iq8[ib].d); + auto scales1 = vcvt_f32_f16(vget_low_f16 (scales16)); + auto scales2 = vcvt_f32_f16(vget_high_f16(scales16)); + for (int j = 0; j < 16; ++j) qx[j] = vld1q_s8(iq8[ib].qs + 16*j); + int32x4_t sumi1, sumi2; + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_0 *)q8.y[iy]; + qx_0_q8_0_dot(qx, qy[ib].qs, sumi1, sumi2); + auto dy = vdupq_n_f32(GGML_FP16_TO_FP32(qy[ib].d)); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], vmulq_f32(scales1, dy), vcvtq_f32_s32(sumi1)); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], vmulq_f32(scales2, dy), vcvtq_f32_s32(sumi2)); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix+0, iy, acc[2*iy+0]); + info.store(ix+4, iy, acc[2*iy+1]); + acc[2*iy] = acc[2*iy+1] = vdupq_n_f32(0.f); + } + } +} + +typedef struct { + ggml_half d[16]; + int8_t qs[256]; +} block_q8_1_r8; + +template +void mul_mat_q8_1_r8_q8_1(int n, const void * vx, size_t bx, const DataInfo& info, int nrc_x) { + GGML_ASSERT(nrc_x%8 == 0); + Q8 q8(info); + int nb = n / QK8_0; + float32x4_t acc[2*nrc_y] = {}; + int8x16_t qx[16]; + float d8[8*nrc_y]; + for (int ix = 0; ix < nrc_x; ix += 8) { + const block_q8_1_r8 * iq8 = (const block_q8_1_r8 *)((const char *)vx + ix*bx); + for (int ib4 = 0; ib4 < nb/4; ++ib4) { + for (int iy = 0; iy < nrc_y; ++iy) { + vst1q_f32(d8+8*iy+0, vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][ib4].d+0))); + vst1q_f32(d8+8*iy+4, vcvt_f32_f16(vld1_f16((const float16_t *)q8.y[iy][ib4].d+4))); + } + for (int k = 0; k < 4; ++k) { + auto scales16 = vld1q_f16((const float16_t *)iq8[4*ib4+k].d); + auto scales1 = vcvt_f32_f16(vget_low_f16 (scales16)); + auto scales2 = vcvt_f32_f16(vget_high_f16(scales16)); + auto m16 = vld1q_f16((const float16_t *)iq8[4*ib4+k].d+8); + auto m1 = vcvt_f32_f16(vget_low_f16 (m16)); + auto m2 = vcvt_f32_f16(vget_high_f16(m16)); + for (int j = 0; j < 16; ++j) qx[j] = vld1q_s8(iq8[4*ib4+k].qs + 16*j); + int32x4_t sumi1, sumi2; + for (int iy = 0; iy < nrc_y; ++iy) { + qx_0_q8_0_dot(qx, q8.y[iy][ib4].qs+32*k, sumi1, sumi2); + auto dy = vdupq_n_f32(d8[8*iy+k]); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], vmulq_f32(scales1, dy), vcvtq_f32_s32(sumi1)); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], vmulq_f32(scales2, dy), vcvtq_f32_s32(sumi2)); + auto my = vdupq_n_f32(d8[8*iy+k+4]); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], m1, my); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], m2, my); + } + } + } + for (int ib = 4*(nb/4); ib < nb; ++ib) { + auto scales16 = vld1q_f16((const float16_t *)iq8[ib].d); + auto scales1 = vcvt_f32_f16(vget_low_f16 (scales16)); + auto scales2 = vcvt_f32_f16(vget_high_f16(scales16)); + auto m16 = vld1q_f16((const float16_t *)iq8[ib].d+8); + auto m1 = vcvt_f32_f16(vget_low_f16 (m16)); + auto m2 = vcvt_f32_f16(vget_high_f16(m16)); + for (int j = 0; j < 16; ++j) qx[j] = vld1q_s8(iq8[ib].qs + 16*j); + int32x4_t sumi1, sumi2; + for (int iy = 0; iy < nrc_y; ++iy) { + auto qy = (const block_q8_1 *)q8.y[iy]; + qx_0_q8_0_dot(qx, qy[ib].qs, sumi1, sumi2); + auto dy = vdupq_n_f32(GGML_FP16_TO_FP32(qy[ib].d)); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], vmulq_f32(scales1, dy), vcvtq_f32_s32(sumi1)); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], vmulq_f32(scales2, dy), vcvtq_f32_s32(sumi2)); + auto my = vdupq_n_f32(GGML_FP16_TO_FP32(qy[ib].s)); + acc[2*iy+0] = vfmaq_f32(acc[2*iy+0], m1, my); + acc[2*iy+1] = vfmaq_f32(acc[2*iy+1], m2, my); + } + } + for (int iy = 0; iy < nrc_y; ++iy) { + info.store(ix+0, iy, acc[2*iy+0]); + info.store(ix+4, iy, acc[2*iy+1]); + acc[2*iy] = acc[2*iy+1] = vdupq_n_f32(0.f); + } + } +} + +struct DeqQ40 { + const int8x16_t m8 = vdupq_n_s8(-8); + const uint8x16_t ml = vdupq_n_s8(0xf); + inline int8x16x2_t dequant(const block_q4_0& x) const { + auto bits = vld1q_u8(x.qs); + return { vaddq_s8(vreinterpretq_s8_u8(vandq_u8(bits, ml)), m8), vaddq_s8(vreinterpretq_s8_u8(vshrq_n_u8(bits, 4)), m8) }; + } +}; + +struct DeqQ41 { + const uint8x16_t ml = vdupq_n_s8(0xf); + inline int8x16x2_t dequant(const block_q4_1& x) const { + auto bits = vld1q_u8(x.qs); + return { vreinterpretq_s8_u8(vandq_u8(bits, ml)), vreinterpretq_s8_u8(vshrq_n_u8(bits, 4)) }; + } +}; + +struct DeqIQ4NL { + const int8x16_t mt = load_values(); + const uint8x16_t ml = vdupq_n_s8(0xf); + inline int8x16x2_t dequant(const block_iq4_nl& x) const { + auto bits = vld1q_u8(x.qs); + return { vqtbl1q_s8(mt, vandq_u8(bits, ml)), vqtbl1q_s8(mt, vshrq_n_u8(bits, 4)) }; + } + static inline int8x16_t load_values() { return vld1q_s8(iq4k_values); } +}; + +struct DeqMXFP4 { + const int8x16_t mt = load_values(); + const uint8x16_t ml = vdupq_n_s8(0xf); + inline int8x16x2_t dequant(const block_mxfp4& x) const { + auto bits = vld1q_u8(x.qs); + return { vqtbl1q_s8(mt, vandq_u8(bits, ml)), vqtbl1q_s8(mt, vshrq_n_u8(bits, 4)) }; + } + static inline int8x16_t load_values() { return vld1q_s8(kvalues_mxfp4); } +}; + +struct DeqQ50 { + + inline int8x16x2_t dequant(const block_q5_0& x) const { + int8x16x2_t r; + bits.prepare1(x.qs, r.val); + auto qh = x.qh; + r.val[0] = vreinterpretq_s8_u8(vorrq_u8(vreinterpretq_u8_s8(r.val[0]), vandq_u8(mh, hbits.to_negated_bytes(qh+0)))); + r.val[1] = vreinterpretq_s8_u8(vorrq_u8(vreinterpretq_u8_s8(r.val[1]), vandq_u8(mh, hbits.to_negated_bytes(qh+2)))); + return r; + } + + Q4LegacyBits bits; + HighBit5Legacy hbits; + const uint8x16_t mh = vdupq_n_u8(0xf0); +}; + +struct DeqQ51 { + + inline int8x16x2_t dequant(const block_q5_1& x) const { + int8x16x2_t r; + bits.prepare1(x.qs, r.val); + auto qh = x.qh; + r.val[0] = vreinterpretq_s8_u8(vorrq_u8(vreinterpretq_u8_s8(r.val[0]), vandq_u8(mh, hbits.to_bytes(qh+0)))); + r.val[1] = vreinterpretq_s8_u8(vorrq_u8(vreinterpretq_u8_s8(r.val[1]), vandq_u8(mh, hbits.to_bytes(qh+2)))); + return r; + } + + Q4LegacyBits bits; + HighBit5Legacy hbits; + const uint8x16_t mh = vdupq_n_u8(0x10); +}; + +struct DeqQ60 { + + inline int8x16x2_t dequant(const block_q6_0& x) const { + int8x16x2_t r; + bits.prepare1(x.qs, r.val); + auto qh8 = vld1_u8(x.qh); + auto qh = vcombine_u8(vshl_n_u8(qh8, 4), qh8); + r.val[0] = vaddq_s8(vorrq_u8(r.val[0], vandq_u8(qh, hmask)), m32); + r.val[1] = vaddq_s8(vorrq_u8(r.val[1], vandq_u8(vshrq_n_u8(qh, 2), hmask)), m32); + return r; + } + + Q4LegacyBits bits; + const int8x16_t m32 = vdupq_n_s8(-32); + const uint8x16_t hmask = vdupq_n_u8(0x30); +}; + +struct DeqQ80 { + inline int8x16x2_t dequant(const block_q8_0& x) const { + return vld1q_s8_x2(x.qs); + } +}; + +template +void iqk_convert_qX_q80_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK4_0 == 0); + GGML_ASSERT(nrc_x%8 == 0); + + const int nb = n/QK8_0; + + block_q8_0_r8 * y = (block_q8_0_r8 *)vy; + + const Block * x8[8]; + + uint32_t block[8]; + + Dequantizer deq; + + for (int ix = 0; ix < nrc_x; ix += 8) { + + for (int k = 0; k < 8; ++k) x8[k] = (const Block *)((const char *)vx + (ix + k)*bx); + + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + if constexpr (std::is_same_v) { + y[i].d[k] = GGML_FP32_TO_FP16(GGML_E8M0_TO_FP32_HALF(x8[k][i].e)); + } else { + y[i].d[k] = x8[k][i].d; + } + vst1q_s8_x2((int8_t *)block, deq.dequant(x8[k][i])); + auto qs = (uint32_t *)y[i].qs; + for (int l = 0; l < 4; ++l) { + qs[8*l + k + 0] = block[l + 0]; + qs[8*l + k + 32] = block[l + 4]; + } + } + } + y += nb; + } +} + +template +void iqk_convert_qX_1_q8_1_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) { + GGML_ASSERT(n%QK4_0 == 0); + GGML_ASSERT(nrc_x%8 == 0); + + const int nb = n/QK8_0; + + block_q8_1_r8 * y = (block_q8_1_r8 *)vy; + + const Block * x8[8]; + + uint32_t block[8]; + + Dequantizer deq; + + for (int ix = 0; ix < nrc_x; ix += 8) { + + for (int k = 0; k < 8; ++k) x8[k] = (const Block *)((const char *)vx + (ix + k)*bx); + + for (int i = 0; i < nb; ++i) { + for (int k = 0; k < 8; ++k) { + y[i].d[k+0] = x8[k][i].d; + y[i].d[k+8] = x8[k][i].m; + vst1q_s8_x2((int8_t *)block, deq.dequant(x8[k][i])); + auto qs = (uint32_t *)y[i].qs; + for (int l = 0; l < 4; ++l) { + qs[8*l + k + 0] = block[l + 0]; + qs[8*l + k + 32] = block[l + 4]; + } + } + } + y += nb; + } +} + +} + +bool iqk_convert_legacy_quants_q8_r8(int type, int n, const void * vx, size_t bx, void * vy, int nrc_x) { + switch (type) { + case GGML_TYPE_Q4_0 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q4_1 : iqk_convert_qX_1_q8_1_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q5_0 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q5_1 : iqk_convert_qX_1_q8_1_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q6_0 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_IQ4_NL: iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_MXFP4 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + case GGML_TYPE_Q8_0 : iqk_convert_qX_q80_r8(n, vx, bx, vy, nrc_x); break; + default: return false; + } + return true; +} + +bool iqk_set_kernels_legacy_quants(int ne00, int typeA, int typeB, std::array& kernels, mul_mat_t& func16) { + + if (ne00%QK8_0 != 0) return false; + + auto etypeA = ggml_type(typeA); + auto expected_typeB = etypeA == GGML_TYPE_Q4_1 || etypeA == GGML_TYPE_Q5_1 || etypeA == GGML_TYPE_Q8_1 ? GGML_TYPE_Q8_1_X4 : GGML_TYPE_Q8_0_X4; + if (ggml_type(typeB) != expected_typeB) return false; + + func16 = nullptr; + + switch (typeA) { + case GGML_TYPE_Q4_0: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_0_q8_0, DequantizerQ40, kernels); + break; + case GGML_TYPE_Q4_1: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_1_q8_1, DequantizerQ41, kernels); + break; + case GGML_TYPE_Q5_0: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_0_q8_0, DequantizerQ50, kernels); + break; + case GGML_TYPE_Q5_1: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_1_q8_1, DequantizerQ51, kernels); + break; + case GGML_TYPE_Q6_0: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_0_q8_0, DequantizerQ60, kernels); + break; + case GGML_TYPE_Q8_0: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_0_q8_0, DequantizerQ80, kernels); + break; + case GGML_TYPE_IQ4_NL: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_0_q8_0, DequantizerIQ4NL, kernels); + break; + case GGML_TYPE_MXFP4: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qX_0_q8_0, DequantizerMXFP4, kernels); + break; + case GGML_TYPE_Q4_0_R8: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qx_r8_q8_0, Q4_0_R8_Dequantizer, kernels); + break; + case GGML_TYPE_Q5_0_R4: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qx_r4_q8_0, Q5_0_R4_Dequantizer, kernels); + break; + case GGML_TYPE_Q6_0_R4: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qx_r4_q8_0, Q6_0_R4_Dequantizer, kernels); + break; + case GGML_TYPE_Q8_0_R8: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_0_r8_q8_0, kernels); + break; + case GGML_TYPE_Q8_1: + IQK_SET_MUL_MAT_FUNCTIONS(mul_mat_q8_1_r8_q8_1, kernels); + break; + case GGML_TYPE_IQ4_NL_R4: + IQK_SET_MUL_MAT_FUNCTIONS_T(mul_mat_qx_r4_q8_0, IQ4_NL_R4_Dequantizer, kernels); + break; + default: + return false; + } + + return true; +} + +#endif + +namespace { +template +inline std::pair mul_mat_kernel(int int_typeA, int nq) { + auto typeA = ggml_type(int_typeA); + constexpr int kMaxQ = 8; +#define MAKE_FUNCS(mul_mat, n) \ + if (n >= kMaxQ) return std::make_pair(mul_mat, kMaxQ>, kMaxQ);\ + else {\ + switch (n) {\ + case 1: return std::make_pair(mul_mat, 1>, 1);\ + case 2: return std::make_pair(mul_mat, 2>, 2);\ + case 3: return std::make_pair(mul_mat, 3>, 3);\ + case 4: return std::make_pair(mul_mat, 4>, 4);\ + case 5: return std::make_pair(mul_mat, 5>, 5);\ + case 6: return std::make_pair(mul_mat, 6>, 6);\ + case 7: return std::make_pair(mul_mat, 7>, 7);\ + }\ + } +#define MAKE_FUNCS2(mul_mat, block, n) \ + if (n >= kMaxQ) return std::make_pair(mul_mat, kMaxQ, block>, kMaxQ);\ + else {\ + switch (n) {\ + case 1: return std::make_pair(mul_mat, 1, block>, 1);\ + case 2: return std::make_pair(mul_mat, 2, block>, 2);\ + case 3: return std::make_pair(mul_mat, 3, block>, 3);\ + case 4: return std::make_pair(mul_mat, 4, block>, 4);\ + case 5: return std::make_pair(mul_mat, 5, block>, 5);\ + case 6: return std::make_pair(mul_mat, 6, block>, 6);\ + case 7: return std::make_pair(mul_mat, 7, block>, 7);\ + }\ + } +#define MAKE_FUNCS_ONLY_NRC(mul_mat, n) \ + if (n >= kMaxQ) return std::make_pair(mul_mat, kMaxQ);\ + else {\ + switch (n) {\ + case 1: return std::make_pair(mul_mat<1>, 1);\ + case 2: return std::make_pair(mul_mat<2>, 2);\ + case 3: return std::make_pair(mul_mat<3>, 3);\ + case 4: return std::make_pair(mul_mat<4>, 4);\ + case 5: return std::make_pair(mul_mat<5>, 5);\ + case 6: return std::make_pair(mul_mat<6>, 6);\ + case 7: return std::make_pair(mul_mat<7>, 7);\ + }\ + } + if (typeA == GGML_TYPE_Q8_0) { +#ifdef __aarch64__ + MAKE_FUNCS(mul_mat_qX_0_q8_0, 1); + if (nq == 2) return std::make_pair(mul_mat_qX_0_q8_2_Tx, 2); + if (nq == 4) return std::make_pair(mul_mat_qX_0_q8_2_Tx, 4); + MAKE_FUNCS(mul_mat_qX_1_q8_2_T, 1); + //if (nq == 2) return std::make_pair(mul_mat_qX_0_q8_0_Tx, 2); + //if (nq == 4) return std::make_pair(mul_mat_qX_0_q8_0_Tx, 4); + if (nq == 1) return std::make_pair(mul_mat_qX_0_q8_0_T, 1); + if (nq == 2) return std::make_pair(mul_mat_qX_0_q8_0_T, 2); + if (nq == 4) return std::make_pair(mul_mat_qX_0_q8_0_T, 4); + if (nq == 3) return std::make_pair(mul_mat_qX_0_q8_0_T, 3); + if (nq == 5) return std::make_pair(mul_mat_qX_0_q8_0_T, 5); + if (nq == 6) return std::make_pair(mul_mat_qX_0_q8_0_T, 6); + if (nq == 7) return std::make_pair(mul_mat_qX_0_q8_0_T, 7); + return std::make_pair(mul_mat_qX_0_q8_0_T, kMaxQ); +#endif +#endif + } + else if (typeA == GGML_TYPE_Q8_0_R8) { +#ifdef __aarch64__ + MAKE_FUNCS_ONLY_NRC(mul_mat_q8_0_r8_q8_0, nq); +#else + MAKE_FUNCS_ONLY_NRC(mul_mat_q8_0_r8_q8_2, nq); +#endif + } + else if (typeA == GGML_TYPE_Q6_0) { +#ifdef __aarch64__ + MAKE_FUNCS(mul_mat_qX_0_q8_0, 1); + if (nq == 2) return std::make_pair(mul_mat_qX_0_q8_2_Tx, 2); + if (nq == 4) return std::make_pair(mul_mat_qX_0_q8_2_Tx, 4); + MAKE_FUNCS(mul_mat_qX_1_q8_2_T, 1); + if (nq == 2) return std::make_pair(mul_mat_qX_0_q8_2_Tx, 2); + if (nq == 4) return std::make_pair(mul_mat_qX_0_q8_2_Tx, 4); + MAKE_FUNCS(mul_mat_qX_1_q8_2_T(nullptr, 0); +} + +inline std::pair mul_mat_kernel(int int_typeA, int nq, int k_step) { + switch (k_step) { + case 32: return mul_mat_kernel< 32>(int_typeA, nq); + case 64: return mul_mat_kernel< 64>(int_typeA, nq); + case 128: return mul_mat_kernel<128>(int_typeA, nq); + default: GGML_ABORT("Fatal error"); + } +} +} + +void iqk_gemm_legacy_fa(int D, int nq, int type_k, const char * k, size_t stride_k, DataInfo& info, int k_step) { + auto [mul_mat, nrc_q] = mul_mat_kernel(type_k, nq, k_step); + for (int iq = 0; iq < nq/nrc_q; ++iq) { + mul_mat(D, k, stride_k, info, k_step); + info.cur_y += nrc_q; + } + int iq = nrc_q*(nq/nrc_q); + if (iq < nq) { + auto [mul_mat1, nrc_q1] = mul_mat_kernel(type_k, nq - iq, k_step); + GGML_ASSERT(nrc_q1 == nq - iq); + mul_mat1(D, k, stride_k, info, k_step); + } +} + +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_gemm_legacy_quants.h b/llama.cpp/ggml/src/iqk/iqk_gemm_legacy_quants.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_gemm_legacy_quants.h @@ -0,0 +1,16 @@ +#pragma once + +#include "iqk_common.h" + +#ifdef IQK_IMPLEMENT + +#include +#include + +bool iqk_set_kernels_legacy_quants(int ne00, int typeA, int typeB, std::array& kernels, mul_mat_t& func16); + +void iqk_gemm_legacy_fa(int D, int nq, int type_k, const char * k, size_t stride_k, DataInfo& info, int k_step); + +bool iqk_convert_legacy_quants_q8_r8(int type, int n, const void * vx, size_t bx, void * vy, int nrc_x); + +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_ggml_type_ext.h b/llama.cpp/ggml/src/iqk/iqk_ggml_type_ext.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_ggml_type_ext.h @@ -0,0 +1,192 @@ +// opencoti F5-opt W2 (#290) — ik_llama ggml_type enum DELTA. +// IK-only enumerators (names absent from llamafile's ggml.h), as +// ((ggml_type)N) macros so iqk code can switch/compare on them. Values +// are >41 (llamafile's max), so they never collide with llamafile's enum. +// VENDORED from ik_llama.cpp @ 8960c5ba ggml/include/ggml.h. MIT, (C) 2024 IK. +#ifndef IQK_GGML_TYPE_EXT_H +#define IQK_GGML_TYPE_EXT_H +#include "ggml.h" +#ifndef GGML_TYPE_Q4_0_4_4 +#define GGML_TYPE_Q4_0_4_4 ((ggml_type)31) +#endif +#ifndef GGML_TYPE_Q4_0_4_8 +#define GGML_TYPE_Q4_0_4_8 ((ggml_type)32) +#endif +#ifndef GGML_TYPE_Q4_0_8_8 +#define GGML_TYPE_Q4_0_8_8 ((ggml_type)33) +#endif +#ifndef GGML_TYPE_I2_S +#define GGML_TYPE_I2_S ((ggml_type)36) +#endif +#ifndef GGML_TYPE_Q1_0_G128 +#define GGML_TYPE_Q1_0_G128 ((ggml_type)41) +#endif +#ifndef GGML_TYPE_Q8_0_X4 +#define GGML_TYPE_Q8_0_X4 ((ggml_type)97) +#endif +#ifndef GGML_TYPE_Q8_1_X4 +#define GGML_TYPE_Q8_1_X4 ((ggml_type)98) +#endif +#ifndef GGML_TYPE_Q8_2_X4 +#define GGML_TYPE_Q8_2_X4 ((ggml_type)99) +#endif +#ifndef GGML_TYPE_Q6_0 +#define GGML_TYPE_Q6_0 ((ggml_type)133) +#endif +#ifndef GGML_TYPE_IQ1_BN +#define GGML_TYPE_IQ1_BN ((ggml_type)134) +#endif +#ifndef GGML_TYPE_IQ2_BN +#define GGML_TYPE_IQ2_BN ((ggml_type)135) +#endif +#ifndef GGML_TYPE_Q8_K64 +#define GGML_TYPE_Q8_K64 ((ggml_type)136) +#endif +#ifndef GGML_TYPE_IQ2_K +#define GGML_TYPE_IQ2_K ((ggml_type)137) +#endif +#ifndef GGML_TYPE_IQ3_K +#define GGML_TYPE_IQ3_K ((ggml_type)138) +#endif +#ifndef GGML_TYPE_IQ4_K +#define GGML_TYPE_IQ4_K ((ggml_type)139) +#endif +#ifndef GGML_TYPE_IQ5_K +#define GGML_TYPE_IQ5_K ((ggml_type)140) +#endif +#ifndef GGML_TYPE_IQ6_K +#define GGML_TYPE_IQ6_K ((ggml_type)141) +#endif +#ifndef GGML_TYPE_IQ4_KS +#define GGML_TYPE_IQ4_KS ((ggml_type)144) +#endif +#ifndef GGML_TYPE_IQ2_KS +#define GGML_TYPE_IQ2_KS ((ggml_type)145) +#endif +#ifndef GGML_TYPE_IQ4_KSS +#define GGML_TYPE_IQ4_KSS ((ggml_type)146) +#endif +#ifndef GGML_TYPE_Q8_K16 +#define GGML_TYPE_Q8_K16 ((ggml_type)147) +#endif +#ifndef GGML_TYPE_Q8_K32 +#define GGML_TYPE_Q8_K32 ((ggml_type)148) +#endif +#ifndef GGML_TYPE_Q8_KR8 +#define GGML_TYPE_Q8_KR8 ((ggml_type)149) +#endif +#ifndef GGML_TYPE_Q8_K128 +#define GGML_TYPE_Q8_K128 ((ggml_type)150) +#endif +#ifndef GGML_TYPE_Q8_KV +#define GGML_TYPE_Q8_KV ((ggml_type)151) +#endif +#ifndef GGML_TYPE_IQ5_KS +#define GGML_TYPE_IQ5_KS ((ggml_type)152) +#endif +#ifndef GGML_TYPE_IQ2_KT +#define GGML_TYPE_IQ2_KT ((ggml_type)153) +#endif +#ifndef GGML_TYPE_IQ3_KT +#define GGML_TYPE_IQ3_KT ((ggml_type)154) +#endif +#ifndef GGML_TYPE_IQ4_KT +#define GGML_TYPE_IQ4_KT ((ggml_type)155) +#endif +#ifndef GGML_TYPE_IQ3_KS +#define GGML_TYPE_IQ3_KS ((ggml_type)156) +#endif +#ifndef GGML_TYPE_IQ2_KL +#define GGML_TYPE_IQ2_KL ((ggml_type)157) +#endif +#ifndef GGML_TYPE_IQ1_KT +#define GGML_TYPE_IQ1_KT ((ggml_type)158) +#endif +#ifndef GGML_TYPE_Q4_0_R8 +#define GGML_TYPE_Q4_0_R8 ((ggml_type)202) +#endif +#ifndef GGML_TYPE_Q5_0_R4 +#define GGML_TYPE_Q5_0_R4 ((ggml_type)206) +#endif +#ifndef GGML_TYPE_Q8_0_R8 +#define GGML_TYPE_Q8_0_R8 ((ggml_type)208) +#endif +#ifndef GGML_TYPE_Q2_K_R4 +#define GGML_TYPE_Q2_K_R4 ((ggml_type)210) +#endif +#ifndef GGML_TYPE_Q3_K_R4 +#define GGML_TYPE_Q3_K_R4 ((ggml_type)211) +#endif +#ifndef GGML_TYPE_Q4_K_R4 +#define GGML_TYPE_Q4_K_R4 ((ggml_type)212) +#endif +#ifndef GGML_TYPE_Q5_K_R4 +#define GGML_TYPE_Q5_K_R4 ((ggml_type)213) +#endif +#ifndef GGML_TYPE_Q6_K_R4 +#define GGML_TYPE_Q6_K_R4 ((ggml_type)214) +#endif +#ifndef GGML_TYPE_IQ2_XXS_R4 +#define GGML_TYPE_IQ2_XXS_R4 ((ggml_type)216) +#endif +#ifndef GGML_TYPE_IQ2_XS_R4 +#define GGML_TYPE_IQ2_XS_R4 ((ggml_type)217) +#endif +#ifndef GGML_TYPE_IQ3_XXS_R4 +#define GGML_TYPE_IQ3_XXS_R4 ((ggml_type)218) +#endif +#ifndef GGML_TYPE_IQ1_S_R4 +#define GGML_TYPE_IQ1_S_R4 ((ggml_type)219) +#endif +#ifndef GGML_TYPE_IQ4_NL_R4 +#define GGML_TYPE_IQ4_NL_R4 ((ggml_type)220) +#endif +#ifndef GGML_TYPE_IQ3_S_R4 +#define GGML_TYPE_IQ3_S_R4 ((ggml_type)221) +#endif +#ifndef GGML_TYPE_IQ2_S_R4 +#define GGML_TYPE_IQ2_S_R4 ((ggml_type)222) +#endif +#ifndef GGML_TYPE_IQ4_XS_R8 +#define GGML_TYPE_IQ4_XS_R8 ((ggml_type)223) +#endif +#ifndef GGML_TYPE_IQ1_M_R4 +#define GGML_TYPE_IQ1_M_R4 ((ggml_type)229) +#endif +#ifndef GGML_TYPE_BF16_R16 +#define GGML_TYPE_BF16_R16 ((ggml_type)230) +#endif +#ifndef GGML_TYPE_Q6_0_R4 +#define GGML_TYPE_Q6_0_R4 ((ggml_type)233) +#endif +#ifndef GGML_TYPE_IQ2_BN_R4 +#define GGML_TYPE_IQ2_BN_R4 ((ggml_type)335) +#endif +#ifndef GGML_TYPE_IQ2_K_R4 +#define GGML_TYPE_IQ2_K_R4 ((ggml_type)337) +#endif +#ifndef GGML_TYPE_IQ3_K_R4 +#define GGML_TYPE_IQ3_K_R4 ((ggml_type)338) +#endif +#ifndef GGML_TYPE_IQ4_K_R4 +#define GGML_TYPE_IQ4_K_R4 ((ggml_type)339) +#endif +#ifndef GGML_TYPE_IQ5_K_R4 +#define GGML_TYPE_IQ5_K_R4 ((ggml_type)340) +#endif +#ifndef GGML_TYPE_IQ4_KS_R4 +#define GGML_TYPE_IQ4_KS_R4 ((ggml_type)344) +#endif +#ifndef GGML_TYPE_IQ5_KS_R4 +#define GGML_TYPE_IQ5_KS_R4 ((ggml_type)352) +#endif +#ifndef GGML_TYPE_Q8_K_R16 +#define GGML_TYPE_Q8_K_R16 ((ggml_type)397) +#endif +#ifndef GGML_TYPE_Q8_KV_R8 +#define GGML_TYPE_Q8_KV_R8 ((ggml_type)398) +#endif +#ifndef GGML_TYPE_Q8_K_R8 +#define GGML_TYPE_Q8_K_R8 ((ggml_type)399) +#endif +#endif // IQK_GGML_TYPE_EXT_H diff --git a/llama.cpp/ggml/src/iqk/iqk_mul_mat.h b/llama.cpp/ggml/src/iqk/iqk_mul_mat.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_mul_mat.h @@ -0,0 +1,83 @@ +// +// Copyright (C) 2024 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// + +#pragma once +#include +#include +#include +#include "iqk_config.h" +#ifdef __cplusplus +extern "C" { +#endif + +IQK_API bool iqk_mul_mat(long Nx, long Ny, long ne00, + int typeA, const void * A, long strideA, + int typeB, const void * B, long strideB, + float * C, long stride_C, int ith, int nth); + +IQK_API bool iqk_mul_mat_4d(long Nx, long Ny, long ne00, + long ne02, long ne03, long ne12, long ne13, + long nb02, long nb03, long nb12, long nb13, long nb2, long nb3, + int typeA, const void * A, long strideA, + int typeB, const void * B, long strideB, + float * C, long stride_C, int ith, int nth); + +IQK_API bool iqk_mul_mat_moe(long Nx, long Ny, long ne00, int ne11, + int typeA, const void * A, long strideA, + int typeB, const void * B, long strideB, + float * C, long nb1, long nb2, const void * vrow_mapping, int ith, int nth); + +IQK_API bool iqk_moe_fused_up_gate(long Nx, long Ny, long ne00, int ne11, int unary_op, + int typeA, const void * Aup, const void * Agate, long strideA, + int typeB, const void * B, long strideB, + const char * up_b, const char * gate_b, + float * C, long nb1, long nb2, const void * vrow_mapping, float limit, int ith, int nth); + +IQK_API int iqk_dequant_type(int type, int Ny); + +struct ggml_tensor; + +IQK_API size_t iqk_fa_work_buffer_size(const struct ggml_tensor * dst, int nthread); + +typedef void (*barrier_t) (void *); + +IQK_API bool iqk_flash_attn_noalibi(int type_q, int type_mask, float max_bias, + int neq3, int neq2, long nbq3, long nbq2, + int nek3, int nek2, long nbk3, long nbk2, + int nev3, int nev2, long nbv3, long nbv2, + int ne2, int ne1, long nb1, + int type_k, // type of k + int type_v, // type of v + int Dk, // K head size + int Dv, // V head size + int nq, // number of columns in q + int nk, // number of rows in k + int stride_q, // distance between q columns in bytes + int stride_k, // distance between k rows in bytes + int stride_v, // distance between v rows in bytes + int stride_m, // distance between mask rows (in bytes + const void * q, // q matrix. + const void * k, // k matrix. Assumed to be fp16, nq x nk elements + const void * v, // v matrix. Assumed to be fp16, nq x nk elements + const void * mask, // mask. If not null, assumed to be fp16. nq x nk elements + const void * sinks, // mask. If not null, assumed to be fp16. nq x nk elements + float scale, // scale applied before softmax + float softcap, // if > 0, a "soft-cap" operation is applied before softmax + float * qkv, // v*softmax(scale*(k*q)) + void * work_buffer, barrier_t barrier, void * barrier_data, + int ith, int nth, int n_swa); + +IQK_API void iqk_topk_moe(int n_experts, int n_experts_used, int nrows, const float * logits, + float * weights, int32_t * ids, int ith, int nth); + +IQK_API bool iqk_fused_delta_net(int head_dim, int n_heads, int gqa_ratio, int repeat_type, int n_tokens, int n_seqs, + size_t vnb1, size_t vnb2, size_t vnb3, + const float * q_data, const float * k_data, const float * v_data, const float * g_data, const float * beta_data, + const float * state_in, float * out_data, float * state_out, float * saved_steps, int state_step_stride, int ith, int nth); + +#ifdef __cplusplus +} +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_quantize.cpp b/llama.cpp/ggml/src/iqk/iqk_quantize.cpp new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_quantize.cpp @@ -0,0 +1,10497 @@ +// +// Copyright (C) 2024 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// + +#if GGML_USE_IQK_MULMAT +#include "iqk_mul_mat.h" +#endif +#include "ggml-quants.h" +#include "ggml-impl.h" +#define GGML_COMMON_IMPL_C +#include "ggml-common.h" +#include "iqk_quantize.h" +#include "iqk_config.h" + +#include "iqk_gemm_ktquants.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +inline int nearest_int(float fval) { + assert(fval <= 4194303.f); + float val = fval + 12582912.f; + int i; memcpy(&i, &val, sizeof(int)); + return (i & 0x007fffff) - 0x00400000; +} + +typedef void (*quantize_func_t)(const float * src, void * qdata, int n_per_row, const float * imatrix, const quantize_user_data * user_data); + +struct QHelper { + QHelper(const float * imatrix, const quantize_user_data * user_data, int n_per_row, int block_size) : + m_imatrix(imatrix), m_user_data(user_data), m_n_per_row(n_per_row), m_block_size(block_size) { + if (m_imatrix) { + m_weight.resize(m_n_per_row); + } + } + const float * row_weights(const float * x) { + constexpr float kEps = 1e-9f; + constexpr float kEps2 = kEps*kEps; + if (!m_imatrix) return m_imatrix; + int nblock = m_n_per_row / m_block_size; + for (int ib = 0; ib < nblock; ++ib) { + auto wb_in = m_imatrix + ib*m_block_size; + auto xb = x + ib*m_block_size; + auto wb = m_weight.data() + ib*m_block_size; + float sumw2 = 0, sumx2 = 0, sumwx = 0; + for (int j = 0; j < m_block_size; ++j) { + wb[j] = wb_in[j]; + sumw2 += wb[j]*wb[j]; + sumx2 += xb[j]*xb[j]; + sumwx += wb[j]*std::abs(xb[j]); + } + if (sumw2 > m_block_size*kEps2 && sumx2 > m_block_size*kEps2 && sumwx > m_block_size*kEps2) continue; + for (int j = 0; j < m_block_size; ++j) { + wb[j] = kEps; + } + } + return m_weight.data(); + } + template + void quantize(int nrows, const float * src, void * dst, int row_size, const Func& qfunc) { + auto cdst = (char *)dst; + for (int row = 0; row < nrows; ++row) { + auto weights = row_weights(src); + qfunc(src, cdst, m_n_per_row, weights, m_user_data); + src += m_n_per_row; + cdst += row_size; + } + } +private: + const float * m_imatrix; + const quantize_user_data * m_user_data; + const int m_n_per_row; + const int m_block_size; + std::vector m_weight; +}; + +template +size_t quantize_repack(ggml_type type, const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data, const Func& q_func, const RepackFunc& repack) { + GGML_ASSERT(nrows%n_repack == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(type, n_per_row); + std::vector qtmp(n_repack*row_size); + QHelper helper(imatrix, user_data, n_per_row, block_size); + char * qrow = (char *)dst; + for (int row = 0; row < nrows; row += n_repack) { + helper.quantize(n_repack, src, qtmp.data(), row_size, q_func); + repack(n_repack, n_per_row, (const Block *)qtmp.data(), (Block_repacked *)qrow, false); + src += n_repack*n_per_row; + qrow += n_repack*row_size; + } + return nrows*row_size; +} + + +float make_qx_quants(int n, int nmax, const float * x, int8_t * L, const float * qw) { + float max = 0; + float amax = 0; + for (int i = 0; i < n; ++i) { + float ax = fabsf(x[i]); + if (ax > amax) { amax = ax; max = x[i]; } + } + if (!amax) { // all zero + for (int i = 0; i < n; ++i) L[i] = 0; + return 0.f; + } + float iscale = -nmax / max; + float sumlx = 0; + float suml2 = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = std::max(-nmax, std::min(nmax-1, l)); + L[i] = l + nmax; + sumlx += qw[i]*x[i]*l; + suml2 += qw[i]*l*l; + } + float scale = suml2 ? sumlx/suml2 : 0.0f; + float best = scale * sumlx; + for (int is = -9; is <= 9; ++is) { + if (is == 0) continue; + iscale = -(nmax + 0.1f*is) / max; + sumlx = suml2 = 0; + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + l = std::max(-nmax, std::min(nmax-1, l)); + sumlx += qw[i]*x[i]*l; + suml2 += qw[i]*l*l; + } + if (suml2 > 0 && sumlx*sumlx > best*suml2) { + for (int i = 0; i < n; ++i) { + int l = nearest_int(iscale * x[i]); + L[i] = nmax + std::max(-nmax, std::min(nmax-1, l)); + } + scale = sumlx/suml2; best = scale*sumlx; + } + } + return scale; +} + +struct IQ1BNQuantizer { + int8_t L[QK_IQ1BN]; + void quantize_one_row_1bn(const float * src, block_iq1_bn * y, int n_per_row, const float * imatrix); + void quantize_one_row_2bn(const float * src, block_iq2_bn * y, int n_per_row, const float * imatrix); + static inline float row_max(int n_per_row, const float * src) { + float max_in_row = 0; + for (int j = 0; j < n_per_row; ++j) { + float ax = fabsf(src[j]); + max_in_row = std::max(max_in_row, ax); + } + return max_in_row; + } + // The Makefile has issues dwaling with this? + //static constexpr uint8_t k_mult[5] = {81, 27, 9, 3, 1}; + static const uint8_t k_mult[5]; +}; + +const uint8_t IQ1BNQuantizer::k_mult[5] = {81, 27, 9, 3, 1}; + +void IQ1BNQuantizer::quantize_one_row_1bn(const float * src, block_iq1_bn * y, int n_per_row, const float * imatrix) { + + static const int k_nb[6] = {1, 3, 9, 27, 81, 243}; + (void)imatrix; + + const int nblock = n_per_row/QK_IQ1BN; + + ggml_half * dptr = (ggml_half *)y; + y = (block_iq1_bn *)(dptr + 1); + + float max = 0; + for (int j = 0; j < n_per_row; ++j) max = std::max(max, fabsf(src[j])); + ggml_half d = GGML_FP32_TO_FP16(max); + std::memcpy(dptr, &d, sizeof(d)); + + float thresh = 0.5f*max; + + for (int ib = 0; ib < nblock; ++ib) { + std::memset(&y[ib], 0, sizeof(block_iq1_bn)); + auto xb = src + ib*QK_IQ1BN; + int v13 = 0; + for (int i16 = 0; i16 < QK_IQ1BN/16; ++i16) { + for (int k = 0; k < 3; ++k) { + int idx = 0; + for (int j = 0; j < 5; ++j) { + float v = xb[16*i16 + 5*k + j]; + int q = fabsf(v) < thresh ? 1 : v < 0 ? 0 : 2; + idx += k_nb[j]*q; + } + idx = (256*idx + k_nb[5] - 1)/k_nb[5]; + y[ib].ql[3*i16 + k] = idx; + } + float v = xb[16*i16 + 15]; + int q = fabsf(v) < thresh ? 1 : v < 0 ? 0 : 2; + v13 += k_nb[i16]*q; + } + y[ib].extra = (256*v13 + k_nb[5] - 1)/k_nb[5]; + } +} + +void IQ1BNQuantizer::quantize_one_row_2bn(const float * src, block_iq2_bn * y, int n_per_row, const float * imatrix) { + + (void)imatrix; + + const int nblock = n_per_row/QK_IQ1BN; + + constexpr int Nj = QK_IQ1BN/4; + + float max = 0; + for (int j = 0; j < n_per_row; ++j) max = std::max(max, fabsf(src[j])); + + float * dptr = (float *)y; + *dptr = max; + y = (block_iq2_bn *)(dptr + 1); + float thresh = 0.5f*max; + + for (int ib = 0; ib < nblock; ++ib) { + auto xb = src + QK_IQ1BN*ib; + for (int j = 0; j < QK_IQ1BN; ++j) { + L[j] = fabsf(xb[j]) < thresh ? 1 : xb[j] < 0 ? 0 : 2; + } + for (int j = 0; j < Nj; ++j) { + y[ib].qs[j] = L[j] | (L[j + Nj] << 2) | (L[j + 2*Nj] << 4) | (L[j + 3*Nj] << 6); + } + } +} + +static inline int num_rows([[maybe_unused]] ggml_type type) { +#ifdef HAVE_FANCY_SIMD + switch (type) { + case GGML_TYPE_Q2_K_R4: + case GGML_TYPE_Q3_K_R4: + case GGML_TYPE_Q6_K_R4: + case GGML_TYPE_IQ2_K_R4: + case GGML_TYPE_IQ3_K_R4: + case GGML_TYPE_IQ4_K_R4: + case GGML_TYPE_IQ5_K_R4: + case GGML_TYPE_IQ4_KS_R4: + case GGML_TYPE_IQ5_KS_R4: + case GGML_TYPE_IQ2_XXS_R4: + case GGML_TYPE_IQ2_XS_R4: + case GGML_TYPE_IQ2_S_R4: + case GGML_TYPE_IQ3_XXS_R4: + case GGML_TYPE_IQ1_S_R4: + case GGML_TYPE_IQ1_M_R4: + case GGML_TYPE_IQ3_S_R4: return 4; + case GGML_TYPE_IQ4_NL_R4: + case GGML_TYPE_Q5_0_R4: + case GGML_TYPE_Q6_0_R4: + case GGML_TYPE_IQ2_BN_R4: + case GGML_TYPE_IQ4_XS_R8: + case GGML_TYPE_Q4_K_R4: + case GGML_TYPE_Q5_K_R4: + case GGML_TYPE_Q8_KV: + case GGML_TYPE_Q8_KV_R8: + case GGML_TYPE_Q8_K_R8: return 8; + case GGML_TYPE_Q4_0_R8: + case GGML_TYPE_Q8_0_R8: + case GGML_TYPE_Q8_1: + case GGML_TYPE_Q8_K_R16: + case GGML_TYPE_BF16_R16: return 16; + default: return 1; + } +#else + switch (type) { + case GGML_TYPE_Q2_K_R4: + case GGML_TYPE_Q3_K_R4: + case GGML_TYPE_Q4_K_R4: + case GGML_TYPE_Q5_K_R4: + case GGML_TYPE_Q6_K_R4: + case GGML_TYPE_Q5_0_R4: + case GGML_TYPE_Q6_0_R4: + case GGML_TYPE_IQ4_NL_R4: + case GGML_TYPE_IQ2_K_R4: + case GGML_TYPE_IQ3_K_R4: + case GGML_TYPE_IQ4_K_R4: + case GGML_TYPE_IQ5_K_R4: + case GGML_TYPE_IQ4_KS_R4: + case GGML_TYPE_IQ5_KS_R4: + case GGML_TYPE_IQ2_XXS_R4: + case GGML_TYPE_IQ2_XS_R4: + case GGML_TYPE_IQ2_S_R4: + case GGML_TYPE_IQ3_XXS_R4: + case GGML_TYPE_IQ3_S_R4: + case GGML_TYPE_IQ1_S_R4: + case GGML_TYPE_IQ1_M_R4: + case GGML_TYPE_IQ2_BN_R4: return 4; + case GGML_TYPE_IQ4_XS_R8: + case GGML_TYPE_Q4_0_R8: + case GGML_TYPE_Q8_0_R8: + case GGML_TYPE_Q8_KV: + case GGML_TYPE_Q8_KV_R8: + case GGML_TYPE_Q8_1: + case GGML_TYPE_Q8_K_R8: return 8; + case GGML_TYPE_Q8_K_R16: + case GGML_TYPE_BF16_R16: return 16; + default: return 1; + } +#endif +} + + +} + +void iqk_quantize_any(int from_type, int to_type, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, + uint64_t nb0, uint64_t nb1, uint64_t nb2, uint64_t nb3, + const void * x, void * y, void * work_buffer, + to_float_t to_float, from_float_t from_float, int ith, int nth) { + auto type_x = ggml_type(from_type); + GGML_ASSERT(ggml_type_size(type_x) == nb0); + auto type_y = ggml_type(to_type); + auto row_size_y = ggml_row_size(type_y, ne0); + auto n_interleaved = num_rows(type_y); + GGML_ASSERT(ne1 % n_interleaved == 0); + int64_t ne1i = ne1/n_interleaved; + int64_t nrows = ne1i*ne2*ne3; + int64_t nrows_per_thread = (nrows + nth - 1)/nth; + int64_t first_row = nrows_per_thread*ith; + if (first_row >= nrows) return; + int64_t last_row = std::min(first_row + nrows_per_thread, nrows); + for (int64_t row = first_row; row < last_row; ++row) { + int64_t i3 = row/(ne1i*ne2); + int64_t i2 = (row - i3*ne1i*ne2)/ne1i; + int64_t i1 = row - i3*ne1i*ne2 - i2*ne1i; + auto cx = (const char *)x + i1*n_interleaved*nb1 + i2*nb2 + i3*nb3; + auto cy = (char *)y + (i3*ne1*ne2 + i2*ne1 + i1*n_interleaved)*row_size_y; + // TODO: special case common types such as f16, q8_0 + // (although the performance gains may be too small to justify the added complexity) + if (type_x != GGML_TYPE_F32) { + to_float((const void *)cx, (float *)work_buffer, ne0*n_interleaved); + from_float((const float *)work_buffer, (void *)cy, ne0*n_interleaved); + } else { + from_float((const float *)cx, (void *)cy, ne0*n_interleaved); + } + } +} + + +size_t quantize_iq1_bn(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data *) { + IQ1BNQuantizer iq1bn; + auto row_size = ggml_row_size(GGML_TYPE_IQ1_BN, n_per_row); + auto qrow = (char *)dst; + for (int row = 0; row < nrows; ++row) { + iq1bn.quantize_one_row_1bn(src + row*n_per_row, (block_iq1_bn *)qrow, n_per_row, imatrix); + qrow += row_size; + } + return nrows*row_size; +} + +void quantize_row_iq1_bn_ref(const float * x, block_iq1_bn * y, int64_t k) { + quantize_iq1_bn(x, y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq1_bn(const float * x, void * y, int64_t k) { + quantize_iq1_bn(x, y, 1, k, nullptr, nullptr); +} + +void dequantize_row_iq1_bn(const block_iq1_bn * x, float * y, int64_t k) { + assert(k%QK_IQ1BN == 0); + int nblock = k / QK_IQ1BN; + + for (int i = 0; i < nblock; ++i) { + uint8_t extra = x[i].extra; + auto ql = x[i].ql; + for (int i16 = 0; i16 < QK_IQ1BN/16; ++i16) { + for (int k = 0; k < 3; ++k) { + for (int j = 0; j < 5; ++j) { + uint8_t v = ql[k]*IQ1BNQuantizer::k_mult[j]; + int8_t vs = ((v + (v >> 1)) >> 7); + *y++ = vs - 1; + } + } + ql += 3; + uint8_t v = extra*IQ1BNQuantizer::k_mult[i16]; + int8_t vs = ((v + (v >> 1)) >> 7); + *y++ = vs - 1; + } + } +} + +size_t quantize_iq2_bn(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + IQ1BNQuantizer iq1bn; + auto row_size = ggml_row_size(GGML_TYPE_IQ2_BN, n_per_row); + auto qrow = (char *)dst; + for (int row = 0; row < nrows; ++row) { + iq1bn.quantize_one_row_2bn(src + row*n_per_row, (block_iq2_bn *)qrow, n_per_row, imatrix); + qrow += row_size; + } + return nrows*row_size; +} + +void quantize_row_iq2_bn_ref(const float * x, block_iq2_bn * y, int64_t k) { + quantize_iq2_bn(x, y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq2_bn(const float * x, void * y, int64_t k) { + quantize_iq2_bn(x, y, 1, k, nullptr, nullptr); +} + +void dequantize_row_iq2_bn(const block_iq2_bn * x, float * y, int64_t k) { + assert(k%QK_IQ1BN == 0); + int nblock = k / QK_IQ1BN; + + auto d1 = 1.f, d2 = 0.25f, d3 = d2*0.25f, d4 = d3*0.25f; + auto m = -1.f; + constexpr int Nj = QK_IQ1BN/4; + for (int i = 0; i < nblock; ++i) { + for (int j = 0; j < Nj; ++j) { + y[j+ 0] = d1*(x[i].qs[j] & 0x03) + m; + y[j+1*Nj] = d2*(x[i].qs[j] & 0x0c) + m; + y[j+2*Nj] = d3*(x[i].qs[j] & 0x30) + m; + y[j+3*Nj] = d4*(x[i].qs[j] & 0xc0) + m; + } + y += QK_IQ1BN; + } +} + +namespace { +inline int8_t iq1bn_dequant(uint8_t q, int i) { + uint8_t v = IQ1BNQuantizer::k_mult[i]*q; + //int8_t vs = (v + (v << 1)) >> 8; + int8_t vs = 3*v >> 8; + return vs - 1; +} +} + +static const int8_t iq1bn_values[1280] = { + -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, 0, -1, -1, -1, 0, 0, -1, -1, -1, 1, 0, + -1, -1, -1, -1, 1, -1, -1, -1, 0, 1, -1, -1, -1, 1, 1, -1, -1, -1, -1, -1, 0, -1, -1, 0, -1, 0, -1, -1, 1, -1, 0, -1, + -1, -1, 0, 0, -1, -1, 0, 0, 0, -1, -1, 1, 0, 0, -1, -1, -1, 1, 0, -1, -1, 0, 1, 0, -1, -1, 1, 1, 0, -1, -1, -1, + -1, 1, -1, -1, 0, 0, 0, 0, 0, 0, -1, 1, -1, -1, 1, -1, 1, -1, -1, -1, 0, 1, -1, -1, 0, 0, 1, -1, -1, 1, 0, 1, + -1, -1, -1, 1, 1, -1, -1, 0, 1, 1, -1, -1, 1, 1, 1, -1, -1, -1, -1, -1, 0, -1, 0, -1, -1, 0, -1, 1, -1, -1, 0, -1, + -1, 0, -1, 0, -1, 0, 0, -1, 0, -1, 1, 0, -1, 0, -1, -1, 1, -1, 0, -1, 0, 1, -1, 0, -1, 1, 1, -1, 0, -1, -1, -1, + 0, 0, -1, 0, -1, 0, 0, -1, 0, 0, 0, 0, 0, 1, -1, 0, 0, -1, -1, 0, 0, 0, -1, 0, 0, 0, 0, -1, 1, 0, 0, 0, + -1, -1, 1, 0, 0, -1, 0, 1, 0, 0, -1, 1, 1, 0, 0, -1, -1, -1, 1, 0, -1, 0, -1, 1, 0, -1, 1, -1, 1, 0, -1, -1, + 0, 1, 0, -1, 0, 0, 1, 0, -1, 1, 0, 1, 0, -1, -1, 1, 1, 0, -1, 0, 1, 1, 0, -1, 1, 1, 1, 0, -1, -1, -1, -1, + 1, -1, 0, -1, -1, 1, -1, 1, -1, -1, 1, -1, 0, 0, 0, 0, 0, -1, 0, -1, 1, -1, 0, 0, -1, 1, -1, 1, 0, -1, 1, -1, + -1, 1, -1, 1, -1, 0, 1, -1, 1, -1, 1, 1, -1, 1, -1, -1, -1, 0, 1, -1, 0, -1, 0, 1, -1, 1, -1, 0, 1, -1, -1, 0, + 0, 1, -1, 0, 0, 0, 1, -1, 1, 0, 0, 1, -1, -1, 1, 0, 1, -1, 0, 1, 0, 1, -1, 1, 1, 0, 1, -1, -1, -1, 1, 1, + -1, 0, -1, 1, 1, -1, 1, -1, 1, 1, -1, 0, 0, 0, 0, 0, -1, 0, 1, 1, -1, 0, 0, 1, 1, -1, 1, 0, 1, 1, -1, -1, + 1, 1, 1, -1, 0, 1, 1, 1, -1, 1, 1, 1, 1, -1, -1, -1, -1, -1, 0, 0, -1, -1, -1, 0, 1, -1, -1, -1, 0, -1, 0, -1, + -1, 0, 0, 0, -1, -1, 0, 1, 0, -1, -1, 0, -1, 1, -1, -1, 0, 0, 1, -1, -1, 0, 1, 1, -1, -1, 0, -1, -1, 0, -1, 0, + 0, -1, 0, -1, 0, 1, -1, 0, -1, 0, -1, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, -1, 0, 1, 0, 0, -1, 0, -1, 1, + 0, -1, 0, 0, 1, 0, -1, 0, 1, 1, 0, -1, 0, -1, -1, 1, -1, 0, 0, -1, 1, -1, 0, 1, -1, 1, -1, 0, -1, 0, 1, -1, + 0, 0, 0, 1, -1, 0, 1, 0, 1, -1, 0, -1, 1, 1, -1, 0, 0, 1, 1, -1, 0, 1, 1, 1, -1, 0, -1, -1, -1, 0, 0, 0, + -1, -1, 0, 0, 1, -1, -1, 0, 0, -1, 0, -1, 0, 0, 0, 0, -1, 0, 0, 0, 0, 0, 0, 0, 1, 0, -1, 0, 0, -1, 1, -1, + 0, 0, 0, 1, -1, 0, 0, 1, 1, -1, 0, 0, -1, -1, 0, 0, 0, 0, -1, 0, 0, 0, 1, -1, 0, 0, 0, -1, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, -1, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, -1, -1, 1, 0, 0, 0, -1, + 1, 0, 0, 1, -1, 1, 0, 0, -1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, -1, 1, 1, 0, + 0, 0, 1, 1, 0, 0, 1, 1, 1, 0, 0, -1, -1, -1, 1, 0, 0, -1, -1, 1, 0, 1, -1, -1, 1, 0, -1, 0, -1, 1, 0, 0, + 0, -1, 1, 0, 1, 0, -1, 1, 0, -1, 1, -1, 1, 0, 0, 1, -1, 1, 0, 1, 1, -1, 1, 0, -1, -1, 0, 1, 0, 0, -1, 0, + 1, 0, 1, -1, 0, 1, 0, -1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, -1, 1, 0, 1, 0, + 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, -1, -1, 1, 1, 0, 0, -1, 1, 1, 0, 1, -1, 1, 1, 0, -1, 0, 1, 1, 0, 0, 0, + 1, 1, 0, 1, 0, 1, 1, 0, -1, 1, 1, 1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1, 0, -1, -1, -1, -1, 1, 0, -1, -1, -1, + 1, 1, -1, -1, -1, 1, -1, 0, -1, -1, 1, 0, 0, -1, -1, 1, 1, 0, -1, -1, 1, -1, 1, -1, -1, 1, 0, 0, 0, 0, 0, 0, + 1, -1, -1, 1, 1, 1, -1, -1, 1, -1, -1, 0, -1, 1, 0, -1, 0, -1, 1, 1, -1, 0, -1, 1, -1, 0, 0, -1, 1, 0, 0, 0, + -1, 1, 1, 0, 0, -1, 1, -1, 1, 0, -1, 1, 0, 1, 0, -1, 1, 1, 1, 0, -1, 1, -1, -1, 1, -1, 1, 0, -1, 1, -1, 1, + 1, -1, 1, -1, 1, -1, 0, 1, -1, 1, 0, 0, 1, -1, 1, 1, 0, 1, -1, 1, -1, 1, 1, -1, 1, 0, 0, 0, 0, 0, 0, 1, + 1, -1, 1, 1, 1, 1, -1, 1, -1, -1, -1, 0, 1, 0, -1, -1, 0, 1, 1, -1, -1, 0, 1, -1, 0, -1, 0, 1, 0, 0, -1, 0, + 1, 1, 0, -1, 0, 1, -1, 1, -1, 0, 1, 0, 1, -1, 0, 1, 1, 1, -1, 0, 1, -1, -1, 0, 0, 1, 0, -1, 0, 0, 1, 1, + -1, 0, 0, 1, -1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 1, -1, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, + 0, 0, 1, 1, 0, 0, 1, -1, -1, 1, 0, 1, 0, -1, 1, 0, 1, 1, -1, 1, 0, 1, -1, 0, 1, 0, 1, 0, 0, 1, 0, 1, + 1, 0, 1, 0, 1, -1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 1, -1, -1, -1, 1, 1, 0, -1, -1, 1, 1, 1, -1, + -1, 1, 1, -1, 0, -1, 1, 1, 0, 0, -1, 1, 1, 1, 0, -1, 1, 1, -1, 1, -1, 1, 1, 0, 1, -1, 1, 1, 1, 1, -1, 1, + 1, 0, 0, 0, 0, 0, -1, -1, 0, 1, 1, 0, -1, 0, 1, 1, 1, -1, 0, 1, 1, -1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, + 0, 0, 1, 1, -1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, -1, -1, 1, 1, 1, 0, -1, 1, 1, 1, 1, -1, 1, + 1, 1, -1, 0, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, -1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, +}; + +void ggml_vec_dot_iq1_bn_q8_K64(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(nrc); + + static_assert(QK_IQ1BN == 64, "This dot product implementation for iq1_bn requires a block size of 64"); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ1_BN, vx, 0, GGML_TYPE_Q8_K64, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + + const block_iq1_bn * x = (const block_iq1_bn *)vx; + + const float * d8 = (const float *)vy; + const int8_t * q8 = (const int8_t *)(d8 + 4); + int nblock = n / QK_IQ1BN; + + int sumi[8] = {}; + int8_t q1[16]; + + for (int ii = 0; ii < nblock; ii += 32) { + int16_t sum16[8] = {}; + int nb = std::min(ii + 32, nblock); + for (int i = ii; i < nb; ++i) { + auto ql = x[i].ql; + const int8_t * extra = iq1bn_values + 5*x[i].extra; + for (int i16 = 0; i16 < QK_IQ1BN/16; ++i16) { + for (int k = 0; k < 3; ++k) { + uint8_t q = *ql++; + const int8_t * vs = iq1bn_values + 5*q; + for (int j = 0; j < 5; ++j) q1[5*k+j] = vs[j]; + } + q1[15] = extra[i16]; + // We collect 8 q8 values per block into each element of sum16 + // => 32 x 8 = 256 values in each loop over i, so this cannot overflow the int16_t range + // (q8 is in -127...127, and hence the sum is in -32512...32512 + for (int j = 0; j < 8; ++j) sum16[j] += q8[2*j+0]*q1[2*j+0] + q8[2*j+1]*q1[2*j+1]; + q8 += 16; + } + } + for (int j = 0; j < 8; ++j) sumi[j] += sum16[j]; + } + + *s = d8[0] * (sumi[0] + sumi[1]) + d8[1] * (sumi[2] + sumi[3]) + d8[2] * (sumi[4] + sumi[5]) + d8[3] * (sumi[6] + sumi[7]); +} + +void vec_dot_iq2_bn_q8_K64(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(nrc); + + static_assert(QK_IQ1BN == 64, "This dot product implementation for iq2_bn requires a block size of 64"); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_BN, vx, 0, GGML_TYPE_Q8_K64, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + + constexpr int Nj = QK_IQ1BN/4; + + const block_iq2_bn * x = (const block_iq2_bn *)vx; + int nblock = n / QK_IQ1BN; + + const float * d = (const float *)vy; + const int8_t * q8 = (const int8_t *)(d + 4); + + int sum[16] = { }; + int sum0[4] = { }; + + for (int i = 0; i < nblock; ++i) { + for (int j = 0; j < Nj/4; ++j) { + for (int l = 0; l < 4; ++l) { + sum[4*j + 0] += q8[4*j + l + 0] * (x[i].qs[4*j+l] & 0x03); + sum[4*j + 1] += q8[4*j + l + 1*Nj] * (x[i].qs[4*j+l] & 0x0c); + sum[4*j + 2] += q8[4*j + l + 2*Nj] * (x[i].qs[4*j+l] & 0x30); + sum[4*j + 3] += q8[4*j + l + 3*Nj] * (x[i].qs[4*j+l] & 0xc0); + sum0[j] += q8[4*j + l] + q8[4*j + l + 1*Nj] + q8[4*j + l + 2*Nj] + q8[4*j + l + 3*Nj]; + } + } + q8 += QK_IQ1BN; + } + + float sumf = 0; + for (int j = 0; j < 4; ++j) { + sumf += d[j] * (sum[4*j + 0] + 0.25f*sum[4*j + 1] + 0.0625*sum[4*j + 2] + 0.015625*sum[4*j + 3] - sum0[j]); + } + *s = sumf; + +} + +void quantize_row_q8_K64_ref(const float * x, block_q8_K64 * y, int64_t k) { + + GGML_ASSERT(k >= 8*QK_IQ1BN); + + float * dptr = (float *)y; + auto qs = (int8_t *)(dptr + 8); +#ifdef __ARM_NEON + static const uint8_t k_shuffle[16] = {0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60}; + auto shuffle = vld1q_u8(k_shuffle); + float32x4_t max[4] = { }; + for (int j = 0; j < k; j += 16) { + for (int i = 0; i < 4; ++i) { + auto val = vld1q_f32(x + j + 4*i); + val = vabsq_f32(val); + max[i] = vmaxq_f32(max[i], val); + } + } + float32x4_t vid[4]; + for (int i = 0; i < 4; ++i) { + dptr[i] = vmaxvq_f32(max[i])/127; + float id = dptr[i] > 0 ? 1/dptr[i] : 0.f; + vid[i] = vdupq_n_f32(id); + } + int8x16x4_t q; + int32x4_t qsum = {}; + const int8x16_t m1 = vdupq_n_s8(1); + for (int j = 0; j < k; j += 16) { + for (int i = 0; i < 4; ++i) { + auto val = vld1q_f32(x + j + 4*i); + val = vmulq_f32(vid[i], val); + auto ival = vcvtnq_s32_f32(val); + q.val[i] = vreinterpretq_s8_s32(ival); + } + auto qi = vqtbl4q_s8(q, shuffle); + qsum = ggml_vdotq_s32(qsum, qi, m1); + vst1q_s8(qs, qi); + qs += 16; + } + auto sumf = vmulq_f32(vld1q_f32(dptr), vcvtq_f32_s32(qsum)); + vst1q_f32(dptr + 4, sumf); +#elif defined __AVX__ + __m128 max[4] = {}; + __m128 sign_bit = _mm_set1_ps(-0.f); + for (int j = 0; j < k; j += 16) { + for (int i = 0; i < 4; ++i) { + auto val = _mm_loadu_ps(x + j + 4*i); + val = _mm_andnot_ps(sign_bit, val); + max[i] = _mm_max_ps(max[i], val); + } + } + __m128 vid[4]; + for (int i = 0; i < 4; ++i) { + max[i] = _mm_max_ps(max[i], _mm_movehl_ps(max[i], max[i])); + max[i] = _mm_max_ss(max[i], _mm_movehdup_ps(max[i])); + float maxi = _mm_cvtss_f32(max[i]); + dptr[i] = maxi/127; + float id = dptr[i] > 0 ? 1/dptr[i] : 0.f; + vid[i] = _mm_set1_ps(id); + } + __m128i q[4]; + __m128i sums = _mm_setzero_si128(); + __m128i m1_8 = _mm_set1_epi8(1); + __m128i m1_16 = _mm_set1_epi16(1); + for (int j = 0; j < k; j += 16) { + for (int i = 0; i < 4; ++i) { + auto val = _mm_loadu_ps(x + j + 4*i); + val = _mm_round_ps(_mm_mul_ps(vid[i], val), _MM_ROUND_NEAREST); + q[i] = _mm_cvtps_epi32(val); + } + auto q1 = _mm_packs_epi32(q[0], q[1]); + auto q2 = _mm_packs_epi32(q[2], q[3]); + auto qi = _mm_packs_epi16(q1, q2); + auto aux = _mm_maddubs_epi16(m1_8, qi); + sums = _mm_add_epi32(sums, _mm_madd_epi16(m1_16, aux)); + _mm_storeu_si128((__m128i *)qs, qi); + qs += 16; + } + auto minus = _mm_mul_ps(_mm_loadu_ps(dptr), _mm_cvtepi32_ps(sums)); + _mm_storeu_ps(dptr + 4, minus); +#else + float aux[4] = {0.f, 0.f, 0.f, 0.f}; + for (int j = 0; j < k; j += 16) { + for (int i = 0; i < 4; ++i) { + for (int l = 0; l < 4; ++l) { + float ax = fabsf(x[j+4*i+l]); + aux[i] = std::max(aux[i], ax); + } + } + } + for (int i = 0; i < 4; ++i) { + dptr[i] = aux[i]/127; + aux[i] = dptr[i] > 0 ? 1/dptr[i] : 0.f; + } + int32_t sum[4] = {}; + for (int j = 0; j < k; j += 16) { + for (int i = 0; i < 4; ++i) { + for (int l = 0; l < 4; ++l) { + qs[j+4*i+l] = nearest_int(aux[i]*x[j+4*i+l]); + sum[i] += qs[j+4*i+l]; + } + } + } + for (int i = 0; i < 4; ++i) dptr[4+i] = dptr[i]*sum[i]; +#endif +} + +void quantize_row_q8_K64(const float * x, void * y, int64_t k) { + quantize_row_q8_K64_ref(x, (block_q8_K64 *)y, k); +} + +void quantize_row_q8_K16(const float * x, void * vy, int64_t nk) { + float * dptr = (float *)vy; + int8_t * qy = (int8_t *)(dptr + 5); + int n64 = nk / 64; +#ifdef z__AVX2__ + __m256 sign_bit = _mm256_set1_ps(-0.f); + __m256 vmax[4] = {}; + __m256 vsum[4] = {}; + for (int i64 = 0; i64 < n64; ++i64) { + for (int k = 0; k < 4; ++k) { + auto v1 = _mm256_loadu_ps(x + 64*i64 + 16*k + 0); + auto v2 = _mm256_loadu_ps(x + 64*i64 + 16*k + 8); + vsum[k] = _mm256_add_ps(vsum[k], _mm256_add_ps(v1, v2)); + v1 = _mm256_andnot_ps(sign_bit, v1); + v2 = _mm256_andnot_ps(sign_bit, v2); + vmax[k] = _mm256_max_ps(vmax[k], _mm256_max_ps(v1, v2)); + } + } + __m256 sum = _mm256_add_ps(_mm256_add_ps(vsum[0], vsum[1]), _mm256_add_ps(vsum[2], vsum[3])); + dptr[4] = hsum_float_8(sum); + for (int k = 0; k < 4; ++k) { + float max = hmax_f32_8(vmax[k]); + dptr[k] = max/127; + vmax[k] = _mm256_set1_ps(dptr[k] > 0 ? 1/dptr[k] : 0.f); + } + __m256i ival[8]; + const __m256i perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7); + for (int i64 = 0; i64 < n64; ++i64) { + for (int k = 0; k < 4; ++k) { + __m256 v0 = _mm256_mul_ps(vmax[k], _mm256_loadu_ps(x + 64*i64 + 16*k + 0)); + __m256 v1 = _mm256_mul_ps(vmax[k], _mm256_loadu_ps(x + 64*i64 + 16*k + 8)); + v0 = _mm256_round_ps(v0, _MM_ROUND_NEAREST); + v1 = _mm256_round_ps(v1, _MM_ROUND_NEAREST); + ival[2*k+0] = _mm256_cvtps_epi32(v0); + ival[2*k+1] = _mm256_cvtps_epi32(v1); + } + for (int k = 0; k < 2; ++k) { + auto i0 = _mm256_packs_epi32(ival[4*k+0], ival[4*k+1]); + auto i1 = _mm256_packs_epi32(ival[4*k+2], ival[4*k+3]); + i0 = _mm256_packs_epi16(i0, i1); + i0 = _mm256_permutevar8x32_epi32(i0, perm); + _mm256_storeu_si256((__m256i *)qy, i0); + qy += 32; + } + } +#elif defined z__ARM_NEON + static const uint8_t k_shuffle[16] = {0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60}; + auto shuffle = vld1q_u8(k_shuffle); + float32x4_t vmax[4] = {}; + float32x4_t vsum[4] = {}; + for (int i64 = 0; i64 < n64; ++i64) { + for (int k = 0; k < 4; ++k) { + auto v = vld1q_f32_x4(x + 64*i64 + 16*k); + vsum[k] = vaddq_f32(vsum[k], vaddq_f32(v.val[0], v.val[1])); + vsum[k] = vaddq_f32(vsum[k], vaddq_f32(v.val[2], v.val[3])); + vmax[k] = vmaxq_f32(vmax[k], vmaxq_f32(vabsq_f32(v.val[0]), vabsq_f32(v.val[1]))); + vmax[k] = vmaxq_f32(vmax[k], vmaxq_f32(vabsq_f32(v.val[2]), vabsq_f32(v.val[3]))); + } + } + dptr[4] = vaddvq_f32(vaddq_f32(vaddq_f32(vsum[0], vsum[1]), vaddq_f32(vsum[2], vsum[3]))); + for (int k = 0; k < 4; ++k) { + float max = vmaxvq_f32(vmax[k]); + dptr[k] = max/127; + vmax[k] = vdupq_n_f32(dptr[k] > 0 ? 1/dptr[k] : 0.f); + } + int8x16x4_t q; + for (int i64 = 0; i64 < n64; ++i64) { + for (int k = 0; k < 4; ++k) { + auto v = vld1q_f32_x4(x + 64*i64 + 16*k); + for (int j = 0; j < 4; ++j) { + q.val[j] = vreinterpretq_s8_s32(vcvtnq_s32_f32(vmulq_f32(vmax[k], v.val[j]))); + } + auto qi = vqtbl4q_s8(q, shuffle); + vst1q_s8(qy, qi); + qy += 16; + } + } +#else + float amax[4] = {0.f, 0.f, 0.f, 0.f}; + for (int i64 = 0; i64 < n64; ++i64) { + for (int k = 0; k < 4; ++k) { + for (int j = 0; j < 16; ++j) { + float ax = std::abs(x[64*i64 + 16*k + j]); + amax[k] = std::max(amax[k], ax); + } + } + } + for (int k = 0; k < 4; ++k) { + dptr[k] = amax[k]/127; + amax[k] = dptr[k] > 0 ? 1/dptr[k] : 0.f; + } + int sumi[4] = {}; + for (int i64 = 0; i64 < n64; ++i64) { + for (int k = 0; k < 4; ++k) { + for (int j = 0; j < 16; ++j) { + int ix = nearest_int(amax[k]*x[64*i64 + 16*k + j]); + sumi[k] += ix; + qy[64*i64 + 16*k + j] = ix; + } + } + } + dptr[4] = dptr[0]*sumi[0] + dptr[1]*sumi[1] + dptr[2]*sumi[2] + dptr[3]*sumi[3]; +#endif +} + +void iqk_quantize_q4_0(const float * x, void * vy, int64_t k) { + const int nb = k / QK4_0; + auto y = (block_q4_0 *)vy; +#ifdef __AVX2__ + static_assert(QK4_0 == 32); + __m256 vx[4], rx[4]; + __m256i ix[4]; + auto v7 = _mm256_set1_ps(7.0f); + auto perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7); + for (int ib = 0; ib < nb; ++ib) { + for (int k = 0; k < 4; ++k) { + vx[k] = _mm256_loadu_ps(x + 8*k); + } + auto max1 = _mm256_max_ps(vx[0], vx[1]); + auto max2 = _mm256_max_ps(vx[2], vx[3]); + auto vmax = _mm256_max_ps(max1, max2); + auto min1 = _mm256_min_ps(vx[0], vx[1]); + auto min2 = _mm256_min_ps(vx[2], vx[3]); + auto vmin = _mm256_min_ps(min1, min2); + float max = hmax_float_8(vmax); + float min = hmin_float_8(vmin); + float amax = std::abs(max); + float amin = std::abs(min); + float d, id; + if (amax > amin) { + d = max / -8; + id = amax > 1e-13f ? 1/d : 0.0f; + } else { + d = min / -8; + id = amin > 1e-13f ? 1/d : 0.0f; + } + auto vid = _mm256_set1_ps(id); + auto vsumqx = _mm256_setzero_ps(); + auto vsumq2 = _mm256_setzero_ps(); + for (int k = 0; k < 4; ++k) { + rx[k] = _mm256_mul_ps(vid, vx[k]); + rx[k] = _mm256_round_ps(rx[k], _MM_ROUND_NEAREST); + rx[k] = _mm256_min_ps(rx[k], v7); + ix[k] = _mm256_cvtps_epi32(rx[k]); + ix[k] = _mm256_add_epi32(ix[k], _mm256_set1_epi32(8)); + auto w = _mm256_mul_ps(vx[k], vx[k]); + auto wr = _mm256_mul_ps(w, rx[k]); + vsumqx = _mm256_fmadd_ps(wr, vx[k], vsumqx); + vsumq2 = _mm256_fmadd_ps(wr, rx[k], vsumq2); + } + auto sumq2 = hsum_float_8(vsumq2); + if (sumq2 > 0) { + auto sumqx = hsum_float_8(vsumqx); + d = sumqx/sumq2; + } + y[ib].d = GGML_FP32_TO_FP16(d); + auto i0 = _mm256_packs_epi32(ix[0], ix[1]); + auto i2 = _mm256_packs_epi32(ix[2], ix[3]); + i0 = _mm256_packs_epi16(i0, i2); + i0 = _mm256_permutevar8x32_epi32(i0, perm); + auto q = _mm_or_si128(_mm256_castsi256_si128(i0), _mm_slli_epi16(_mm256_extracti128_si256(i0, 1), 4)); + _mm_storeu_si128((__m128i *)y[ib].qs, q); + x += QK4_0; + } +#else + for (int ib = 0; ib < nb; ++ib) { + float max = 0, amax = 0; + for (int j = 0; j < QK4_0; ++j) { + float ax = std::abs(x[j]); + if (ax > amax) { + amax = ax; max = x[j]; + } + } + float d = max / -8; + float id = amax > 1e-13f ? 1/d : 0.0f; + float sumqx = 0, sumq2 = 0; + for (int j = 0; j < QK4_0/2; ++j) { + float v0 = x[j], v1 = x[j+QK4_0/2]; + int i0 = nearest_int(id*v0), i1 = nearest_int(id*v1); + i0 = std::min(i0, 7); + i1 = std::min(i1, 7); + float w0 = v0*v0, w1 = v1*v1; + sumqx += w0*i0*v0 + w1*i1*v1; + sumq2 += w0*i0*i0 + w1*i1*i1; + y[ib].qs[j] = (i0 + 8) | ((i1 + 8) << 4); + } + if (sumq2 > 0) d = sumqx/sumq2; + y[ib].d = GGML_FP32_TO_FP16(d); + x += QK4_0; + } +#endif +} + +void quantize_row_q8_0_x4(const float * x, void * vy, int64_t k) { + const int nb = k / QK8_0; + const int nb4 = 4*(nb/4); + + block_q8_0 * y = (block_q8_0 *)vy; + block_q8_0_x4 * y4 = (block_q8_0_x4 *)vy; +#if defined(__aarch64__) + for (int i = 0; i < nb; i++) { + int i4 = i/4, ir = i%4; + float32x4_t srcv [8]; + float32x4_t asrcv[8]; + float32x4_t amaxv[8]; + + for (int j = 0; j < 8; j++) srcv[j] = vld1q_f32(x + i*32 + 4*j); + for (int j = 0; j < 8; j++) asrcv[j] = vabsq_f32(srcv[j]); + + for (int j = 0; j < 4; j++) amaxv[2*j] = vmaxq_f32(asrcv[2*j], asrcv[2*j+1]); + for (int j = 0; j < 2; j++) amaxv[4*j] = vmaxq_f32(amaxv[4*j], amaxv[4*j+2]); + for (int j = 0; j < 1; j++) amaxv[8*j] = vmaxq_f32(amaxv[8*j], amaxv[8*j+4]); + + const float amax = vmaxvq_f32(amaxv[0]); + + const float d = amax / ((1 << 7) - 1); + const float id = d ? 1.0f/d : 0.0f; + + if (i < nb4) { + y4[i4].d[ir] = GGML_FP32_TO_FP16(d); + } else { + y[i].d = GGML_FP32_TO_FP16(d); + } + + for (int j = 0; j < 8; j++) { + const float32x4_t v = vmulq_n_f32(srcv[j], id); + const int32x4_t vi = vcvtnq_s32_f32(v); + + if (i < nb4) { + y4[i4].qs[32*ir + 4*j + 0] = vgetq_lane_s32(vi, 0); + y4[i4].qs[32*ir + 4*j + 1] = vgetq_lane_s32(vi, 1); + y4[i4].qs[32*ir + 4*j + 2] = vgetq_lane_s32(vi, 2); + y4[i4].qs[32*ir + 4*j + 3] = vgetq_lane_s32(vi, 3); + } else { + y[i].qs[4*j + 0] = vgetq_lane_s32(vi, 0); + y[i].qs[4*j + 1] = vgetq_lane_s32(vi, 1); + y[i].qs[4*j + 2] = vgetq_lane_s32(vi, 2); + y[i].qs[4*j + 3] = vgetq_lane_s32(vi, 3); + } + } + } +#else + for (int i = 0; i < nb; i++) { + int i4 = i/4, ir = i%4; + // Load elements into 4 AVX vectors + __m256 v0 = _mm256_loadu_ps( x ); + __m256 v1 = _mm256_loadu_ps( x + 8 ); + __m256 v2 = _mm256_loadu_ps( x + 16 ); + __m256 v3 = _mm256_loadu_ps( x + 24 ); + x += 32; + + const __m256 signBit = _mm256_set1_ps( -0.0f ); + __m256 maxAbs = _mm256_andnot_ps( signBit, v0 ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v1 ) ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v2 ) ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v3 ) ); + + __m128 max4 = _mm_max_ps( _mm256_extractf128_ps( maxAbs, 1 ), _mm256_castps256_ps128( maxAbs ) ); + max4 = _mm_max_ps( max4, _mm_movehl_ps( max4, max4 ) ); + max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4 ) ); + const float maxScalar = _mm_cvtss_f32( max4 ); + + const float d = maxScalar / 127.f; + if (i < nb4) { + y4[i4].d[ir] = GGML_FP32_TO_FP16(d); + } else { + y[i].d = GGML_FP32_TO_FP16(d); + } + const float id = ( maxScalar != 0.0f ) ? 127.f / maxScalar : 0.0f; + const __m256 mul = _mm256_set1_ps( id ); + + v0 = _mm256_mul_ps( v0, mul ); + v1 = _mm256_mul_ps( v1, mul ); + v2 = _mm256_mul_ps( v2, mul ); + v3 = _mm256_mul_ps( v3, mul ); + + v0 = _mm256_round_ps( v0, _MM_ROUND_NEAREST ); + v1 = _mm256_round_ps( v1, _MM_ROUND_NEAREST ); + v2 = _mm256_round_ps( v2, _MM_ROUND_NEAREST ); + v3 = _mm256_round_ps( v3, _MM_ROUND_NEAREST ); + + __m256i i0 = _mm256_cvtps_epi32( v0 ); + __m256i i1 = _mm256_cvtps_epi32( v1 ); + __m256i i2 = _mm256_cvtps_epi32( v2 ); + __m256i i3 = _mm256_cvtps_epi32( v3 ); + + // Convert int32 to int16 + i0 = _mm256_packs_epi32( i0, i1 ); // 0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 12, 13, 14, 15 + i2 = _mm256_packs_epi32( i2, i3 ); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31 + // Convert int16 to int8 + i0 = _mm256_packs_epi16( i0, i2 ); // 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27, 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31 + + // We got our precious signed bytes, but the order is now wrong + // These AVX2 pack instructions process 16-byte pieces independently + // The following instruction is fixing the order + const __m256i perm = _mm256_setr_epi32( 0, 4, 1, 5, 2, 6, 3, 7 ); + i0 = _mm256_permutevar8x32_epi32( i0, perm ); + + if (i < nb4) { + _mm256_storeu_si256((__m256i *)y4[i4].qs + ir, i0); + } else { + _mm256_storeu_si256((__m256i *)y[i].qs, i0); + } + } +#endif +} + +namespace { +template +void quantize_row_q8_1_x4_T(const float * x, Block * y, int64_t k) { + assert(k % QK8_1 == 0); + const int nb = k / QK8_1; + + const int nb4 = 4*(nb/4); + Block_x4 * y4 = (Block_x4 *)y; +#if defined(__aarch64__) + for (int i = 0; i < nb; i++) { + int i4 = i/4, ir = i%4; + float32x4_t srcv [8]; + float32x4_t asrcv[8]; + float32x4_t amaxv[8]; + + for (int j = 0; j < 8; j++) srcv[j] = vld1q_f32(x + i*32 + 4*j); + for (int j = 0; j < 8; j++) asrcv[j] = vabsq_f32(srcv[j]); + + for (int j = 0; j < 4; j++) amaxv[2*j] = vmaxq_f32(asrcv[2*j], asrcv[2*j+1]); + for (int j = 0; j < 2; j++) amaxv[4*j] = vmaxq_f32(amaxv[4*j], amaxv[4*j+2]); + for (int j = 0; j < 1; j++) amaxv[8*j] = vmaxq_f32(amaxv[8*j], amaxv[8*j+4]); + + const float amax = vmaxvq_f32(amaxv[0]); + + const float d = amax / ((1 << 7) - 1); + const float id = d ? 1.0f/d : 0.0f; + + if (i < nb4) { + y4[i4].d[ir] = GGML_FP32_TO_FP16(d); + } else { + y[i].d = GGML_FP32_TO_FP16(d); + } + + int32x4_t accv = vdupq_n_s32(0); + + for (int j = 0; j < 8; j++) { + const float32x4_t v = vmulq_n_f32(srcv[j], id); + const int32x4_t vi = vcvtnq_s32_f32(v); + + if (i < nb4) { + y4[i4].qs[QK8_1*ir + 4*j + 0] = vgetq_lane_s32(vi, 0); + y4[i4].qs[QK8_1*ir + 4*j + 1] = vgetq_lane_s32(vi, 1); + y4[i4].qs[QK8_1*ir + 4*j + 2] = vgetq_lane_s32(vi, 2); + y4[i4].qs[QK8_1*ir + 4*j + 3] = vgetq_lane_s32(vi, 3); + } else { + y[i].qs[4*j + 0] = vgetq_lane_s32(vi, 0); + y[i].qs[4*j + 1] = vgetq_lane_s32(vi, 1); + y[i].qs[4*j + 2] = vgetq_lane_s32(vi, 2); + y[i].qs[4*j + 3] = vgetq_lane_s32(vi, 3); + } + + accv = vaddq_s32(accv, vi); + } + + if constexpr (std::is_same_v) { + if (i < nb4) { + y4[i4].d[ir+4] = GGML_FP32_TO_FP16(d * vaddvq_s32(accv)); + } else { + y[i].s = GGML_FP32_TO_FP16(d * vaddvq_s32(accv)); + } + } else { + if (i < nb4) { + y4[i4].d[ir+4] = GGML_FP32_TO_BF16(d * vaddvq_s32(accv)).bits; + } else { + y[i].s = GGML_FP32_TO_BF16(d * vaddvq_s32(accv)).bits; + } + } + } +#else + for (int i = 0; i < nb; i++) { + int i4 = i/4, ir = i%4; + // Load elements into 4 AVX vectors + __m256 v0 = _mm256_loadu_ps( x ); + __m256 v1 = _mm256_loadu_ps( x + 8 ); + __m256 v2 = _mm256_loadu_ps( x + 16 ); + __m256 v3 = _mm256_loadu_ps( x + 24 ); + x += 32; + + // Compute max(abs(e)) for the block + const __m256 signBit = _mm256_set1_ps( -0.0f ); + __m256 maxAbs = _mm256_andnot_ps( signBit, v0 ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v1 ) ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v2 ) ); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps( signBit, v3 ) ); + + __m128 max4 = _mm_max_ps( _mm256_extractf128_ps( maxAbs, 1 ), _mm256_castps256_ps128( maxAbs ) ); + max4 = _mm_max_ps( max4, _mm_movehl_ps( max4, max4 ) ); + max4 = _mm_max_ss( max4, _mm_movehdup_ps( max4 ) ); + const float max_scalar = _mm_cvtss_f32( max4 ); + + // Quantize these floats + float d = max_scalar / 127.f; + if constexpr (std::is_same_v) { + if (i < nb4) { + y4[i4].d[ir] = GGML_FP32_TO_FP16(d); + } else { + y[i].d = GGML_FP32_TO_FP16(d); + } + } else { + auto t = GGML_FP32_TO_BF16(d); + d = ggml_bf16_to_fp32(t); + if (i < nb4) { + y4[i4].d[ir] = t.bits; + } else { + y[i].d = t.bits; + } + } + const float id = d > 0 ? 1/d : 0.f; + const __m256 mul = _mm256_set1_ps( id ); + + // Apply the multiplier + v0 = _mm256_mul_ps( v0, mul ); + v1 = _mm256_mul_ps( v1, mul ); + v2 = _mm256_mul_ps( v2, mul ); + v3 = _mm256_mul_ps( v3, mul ); + + // Round to nearest integer + v0 = _mm256_round_ps( v0, _MM_ROUND_NEAREST ); + v1 = _mm256_round_ps( v1, _MM_ROUND_NEAREST ); + v2 = _mm256_round_ps( v2, _MM_ROUND_NEAREST ); + v3 = _mm256_round_ps( v3, _MM_ROUND_NEAREST ); + + // Convert floats to integers + __m256i i0 = _mm256_cvtps_epi32( v0 ); + __m256i i1 = _mm256_cvtps_epi32( v1 ); + __m256i i2 = _mm256_cvtps_epi32( v2 ); + __m256i i3 = _mm256_cvtps_epi32( v3 ); + + // Compute the sum of the quants and set y[i].s + int isum = hsum_i32_8(_mm256_add_epi32(_mm256_add_epi32(i0, i1), _mm256_add_epi32(i2, i3))); + if constexpr (std::is_same_v) { + if (i < nb4) { + y4[i4].d[ir+4] = GGML_FP32_TO_FP16(d * isum); + } else { + y[i].s = GGML_FP32_TO_FP16(d * isum); + } + } else { + if (i < nb4) { + auto i16 = (int16_t *)y4[i4].d; + i16[ir+4] = isum; + } else { + auto i16 = (int16_t *)&y[i].s; + i16[0] = isum; + } + } + + // Convert int32 to int16 + i0 = _mm256_packs_epi32( i0, i1 ); // 0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 12, 13, 14, 15 + i2 = _mm256_packs_epi32( i2, i3 ); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31 + // Convert int16 to int8 + i0 = _mm256_packs_epi16( i0, i2 ); // 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27, 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31 + + // We got our precious signed bytes, but the order is now wrong + // These AVX2 pack instructions process 16-byte pieces independently + // The following instruction is fixing the order + const __m256i perm = _mm256_setr_epi32( 0, 4, 1, 5, 2, 6, 3, 7 ); + i0 = _mm256_permutevar8x32_epi32( i0, perm ); + + if (i < nb4) { + _mm256_storeu_si256((__m256i *)y4[i4].qs + ir, i0); + } else { + _mm256_storeu_si256((__m256i *)y[i].qs, i0); + } + } +#endif +} +} + +void quantize_row_q8_1_x4(const float * x, void * vy, int64_t k) { + quantize_row_q8_1_x4_T(x, (block_q8_1 *)vy, k); +} + +void quantize_row_q8_2_x4(const float * x, void * vy, int64_t k) { + quantize_row_q8_1_x4_T(x, (block_q8_2 *)vy, k); +} + +// +// ============================================== iq2_K +// + +namespace { + +inline int best_index_iq2nl(const int8_t * values, float x) { + int idx = x < values[1] ? 0 : x > values[2] ? 2 : 1; + return x - values[idx] < values[idx+1] - x ? idx : idx + 1; +} + +void quantize_row_iq2_k_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, + [[maybe_unused]] const quantize_user_data * user_data) { + + constexpr int kBlockSize = 16; + + block_iq2_k * y = (block_iq2_k *)vy; + + float scales[QK_K/kBlockSize]; + float weight[kBlockSize]; + float sumx[kBlockSize+1], sumw[kBlockSize+1]; + float sw[QK_K/kBlockSize]; + int8_t Ls[QK_K/kBlockSize]; + + std::array, kBlockSize> pairs; + + const int8_t * shifted_values = iq2nl_values + 4; + + for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq2_k)); + y[ibl].d = GGML_FP32_TO_FP16(0.f); + + const float * xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 1.5f*sumx2/QK_K; + + uint16_t extra = 0; + + float max_abs_scale = 0; + + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + const float * xb = xbl + kBlockSize*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + sw[ib] = 0; + float amax = 0; + for (int j = 0; j < kBlockSize; ++j) { + sw[ib] += weight[j]; + pairs[j] = {xb[j], j}; + float ax = std::abs(xb[j]); + amax = std::max(amax, ax); + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + std::sort(pairs.begin(), pairs.end()); + sumx[0] = sumw[0] = 0; + for (int j = 0; j < kBlockSize; ++j) { + int jj = pairs[j].second; + sumw[j+1] = sumw[j] + weight[jj]; + sumx[j+1] = sumx[j] + weight[jj]*xb[jj]; + } + float best = 0, d = 0; + bool is_shifted = false; + float sumqx, sumq2; + for (int i1 = 0; i1 < kBlockSize; ++i1) { + for (int i2 = i1; i2 < kBlockSize; ++i2) { + for (int i3 = i2; i3 < kBlockSize; ++i3) { + sumqx = (sumx[i1] - sumx[ 0])*iq2nl_values[0] + (sumx[i2] - sumx[i1])*iq2nl_values[1] + + (sumx[i3] - sumx[i2])*iq2nl_values[2] + (sumx[kBlockSize] - sumx[i3])*iq2nl_values[3]; + sumq2 = (sumw[i1] - sumw[ 0])*iq2nl_values[0]*iq2nl_values[0] + (sumw[i2] - sumw[i1])*iq2nl_values[1]*iq2nl_values[1] + + (sumw[i3] - sumw[i2])*iq2nl_values[2]*iq2nl_values[2] + (sumw[kBlockSize] - sumw[i3])*iq2nl_values[3]*iq2nl_values[3]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; is_shifted = false; + } + sumqx = (sumx[i1] - sumx[ 0])*shifted_values[0] + (sumx[i2] - sumx[i1])*shifted_values[1] + + (sumx[i3] - sumx[i2])*shifted_values[2] + (sumx[kBlockSize] - sumx[i3])*shifted_values[3]; + sumq2 = (sumw[i1] - sumw[ 0])*shifted_values[0]*shifted_values[0] + (sumw[i2] - sumw[i1])*shifted_values[1]*shifted_values[1] + + (sumw[i3] - sumw[i2])*shifted_values[2]*shifted_values[2] + (sumw[kBlockSize] - sumw[i3])*shifted_values[3]*shifted_values[3]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; is_shifted = true; + } + sumqx = (sumx[i1] - sumx[ 0])*iq2nl_values[3] + (sumx[i2] - sumx[i1])*iq2nl_values[2] + + (sumx[i3] - sumx[i2])*iq2nl_values[1] + (sumx[kBlockSize] - sumx[i3])*iq2nl_values[0]; + sumq2 = (sumw[i1] - sumw[ 0])*iq2nl_values[3]*iq2nl_values[3] + (sumw[i2] - sumw[i1])*iq2nl_values[2]*iq2nl_values[2] + + (sumw[i3] - sumw[i2])*iq2nl_values[1]*iq2nl_values[1] + (sumw[kBlockSize] - sumw[i3])*iq2nl_values[0]*iq2nl_values[0]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; is_shifted = false; + } + sumqx = (sumx[i1] - sumx[ 0])*shifted_values[3] + (sumx[i2] - sumx[i1])*shifted_values[2] + + (sumx[i3] - sumx[i2])*shifted_values[1] + (sumx[kBlockSize] - sumx[i3])*shifted_values[0]; + sumq2 = (sumw[i1] - sumw[ 0])*shifted_values[3]*shifted_values[3] + (sumw[i2] - sumw[i1])*shifted_values[2]*shifted_values[2] + + (sumw[i3] - sumw[i2])*shifted_values[1]*shifted_values[1] + (sumw[kBlockSize] - sumw[i3])*shifted_values[0]*shifted_values[0]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; is_shifted = true; + } + } + } + } + scales[ib] = d; + if (is_shifted) extra |= (1 << ib); + + float abs_scale = fabsf(scales[ib]); + max_abs_scale = std::max(max_abs_scale, abs_scale); + } + + if (!max_abs_scale) continue; + float d = make_qx_quants(QK_K/kBlockSize, 8, scales, Ls, sw); + if (!d) continue; + + //float d = -max_scale/8; + y[ibl].extra = extra; + float id = 1/d; + + float sumqx = 0, sumq2 = 0; + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + int ls = nearest_int(id*scales[ib]); + ls = std::max(-8, std::min(7, ls)); + y[ibl].scales[ib/2] |= ((ls + 8) << 4*(ib%2)); + float dl = d * ls; + if (dl) { + const int8_t * block_values = y[ibl].extra & (1 << ib) ? shifted_values : iq2nl_values; + const float * xb = xbl + kBlockSize*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float idl = 1/dl; + int ib32 = ib/2; + int offset = 16*(ib%2); + uint8_t * qs = y[ibl].qs + 32*(ib32/4) + offset; + for (int j = 0; j < 16; ++j) { + const float al = idl*xb[j]; + int ibest = best_index_iq2nl(block_values, al); + qs[j] |= (ibest << 2*(ib32%4)); + float w = weight[j]; + float q = block_values[ibest]*ls; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + } + } + y[ibl].d = GGML_FP32_TO_FP16(1.030f*(sumq2 > 0 ? sumqx/sumq2 : d)); + + } +} +} + +void quantize_row_iq2_k_ref(const float * x, block_iq2_k * y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq2_k(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq2_k(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + block_iq2_k * y = (block_iq2_k *)vy; + quantize_row_iq2_k_ref(x, y, k); +} + +size_t quantize_iq2_k(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + QHelper helper(imatrix, user_data, n_per_row, 16); + auto row_size = ggml_row_size(GGML_TYPE_IQ2_K, n_per_row); + helper.quantize(nrows, src, dst, row_size, quantize_row_iq2_k_impl); + return nrows * row_size; +} + +void dequantize_row_iq2_k(const block_iq2_k * x, float * y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const uint8_t * qs = x[i].qs; + + uint16_t extra = x[i].extra; + + int shift = 0; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + float dl1 = d * ((x[i].scales[ib32] & 0xf) - 8); + float dl2 = d * ((x[i].scales[ib32] >> 4) - 8); + const int8_t * values1 = extra & 1 ? iq2nl_values + 4 : iq2nl_values; + const int8_t * values2 = extra & 2 ? iq2nl_values + 4 : iq2nl_values; + extra >>= 2; + for (int j = 0; j < 16; ++j) { + y[j+ 0] = dl1 * values1[(qs[j+ 0] >> shift) & 3]; + y[j+16] = dl2 * values2[(qs[j+16] >> shift) & 3]; + } + y += 32; + shift += 2; + if (shift == 8) { qs += 32; shift = 0; } + } + + } + +} + +void vec_dot_iq2_k_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_K, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + + GGML_ABORT("not implemented"); + +} + +namespace { +#if defined(__AVX2__) +inline void to_values_i32(__m256i idx, __m256i ivalues, __m256i * iv) { + auto ival = _mm256_shuffle_epi8(ivalues, idx); + auto ival_1 = _mm256_srli_si256(ival, 8); + iv[0] = _mm256_cvtepi8_epi32(_mm256_castsi256_si128(ival)); + iv[1] = _mm256_cvtepi8_epi32(_mm256_castsi256_si128(ival_1)); + iv[2] = _mm256_cvtepi8_epi32(_mm256_extracti128_si256(ival, 1)); + iv[3] = _mm256_cvtepi8_epi32(_mm256_extracti128_si256(ival_1, 1)); +} +inline __m256i to_int8(const __m256i * ibest) { + auto i0 = _mm256_packs_epi32(ibest[0], ibest[1]); // 0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 6, 7, 12, 13, 14, 15 + auto i1 = _mm256_packs_epi32(ibest[2], ibest[3]); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31 + auto idx = _mm256_packs_epi16(i0, i1); // 0, 1, 2, 3, 8, 9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27, 4, 5, 6, 7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31 + auto idx_l = _mm256_castsi256_si128(idx); + auto idx_h = _mm256_extracti128_si256(idx, 1); + auto idx1 = _mm_unpacklo_epi32(idx_l, idx_h); // 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15 + auto idx2 = _mm_unpackhi_epi32(idx_l, idx_h); + return MM256_SET_M128I(idx2, idx1); +} +bool compute_1block_iq2ks(float d, const __m256 * vx, const __m256 * vw, const int8_t * values, __m256i & this_idx, float & best_d, float & score) { + constexpr int kBlockSize = 32; + uint32_t aux32; + std::memcpy(&aux32, values, sizeof(aux32)); + auto ivalues = _mm256_set1_epi32(aux32); + __m256 vbest[8]; + __m256i ibest[8]; + auto val = _mm256_set1_ps(d*values[0]); + auto ival = _mm256_set1_epi32(0); + for (int k = 0; k < kBlockSize/8; ++k) { + auto diff = _mm256_sub_ps(vx[k], val); + vbest[k] = _mm256_mul_ps(diff, diff); + ibest[k] = ival; + diff = _mm256_add_ps(vx[k], val); + vbest[k+4] = _mm256_mul_ps(diff, diff); + ibest[k+4] = ival; + } + for (int j = 1; j < 4; ++j) { + val = _mm256_set1_ps(d*values[j]); + ival = _mm256_set1_epi32(j); + for (int k = 0; k < kBlockSize/8; ++k) { + auto diff = _mm256_sub_ps(vx[k], val); + diff = _mm256_mul_ps(diff, diff); + auto mask = _mm256_cmp_ps(diff, vbest[k], _CMP_LT_OQ); + vbest[k] = _mm256_or_ps(_mm256_and_ps(mask, diff), _mm256_andnot_ps(mask, vbest[k])); + auto imask = _mm256_castps_si256(mask); + ibest[k] = _mm256_or_si256(_mm256_and_si256(imask, ival), _mm256_andnot_si256(imask, ibest[k])); + diff = _mm256_add_ps(vx[k], val); + diff = _mm256_mul_ps(diff, diff); + mask = _mm256_cmp_ps(diff, vbest[k+4], _CMP_LT_OQ); + vbest[k+4] = _mm256_or_ps(_mm256_and_ps(mask, diff), _mm256_andnot_ps(mask, vbest[k+4])); + imask = _mm256_castps_si256(mask); + ibest[k+4] = _mm256_or_si256(_mm256_and_si256(imask, ival), _mm256_andnot_si256(imask, ibest[k+4])); + } + } + bool result = false; + auto idx1 = to_int8(ibest+0); + auto idx2 = to_int8(ibest+4); + to_values_i32(idx1, ivalues, ibest+0); + to_values_i32(idx2, ivalues, ibest+4); + auto vsqx_1 = _mm256_setzero_ps(); + auto vsq2_1 = _mm256_setzero_ps(); + auto vsqx_2 = _mm256_setzero_ps(); + auto vsq2_2 = _mm256_setzero_ps(); + for (int k = 0; k < 4; ++k) { + auto vq1 = _mm256_cvtepi32_ps(ibest[k+0]); + auto vwq1 = _mm256_mul_ps(vw[k], vq1); + auto vq2 = _mm256_cvtepi32_ps(ibest[k+4]); + auto vwq2 = _mm256_mul_ps(vw[k], vq2); + vsqx_1 = _mm256_fmadd_ps(vwq1, vx[k], vsqx_1); + vsq2_1 = _mm256_fmadd_ps(vwq1, vq1, vsq2_1); + vsqx_2 = _mm256_fmadd_ps(vwq2, vx[k], vsqx_2); + vsq2_2 = _mm256_fmadd_ps(vwq2, vq2, vsq2_2); + } + auto sumqx_1 = hsum_float_8(vsqx_1); + auto sumq2_1 = hsum_float_8(vsq2_1); + auto sumqx_2 = hsum_float_8(vsqx_2); + auto sumq2_2 = hsum_float_8(vsq2_2); + if (sumq2_1 > 0) { + best_d = sumqx_1/sumq2_1; + score = sumqx_1 * best_d; + this_idx = idx1; + result = true; + } + if (sumq2_2 > 0 && (!result || sumqx_2*sumqx_2 > score*sumq2_2)) { + best_d = sumqx_2/sumq2_2; + score = sumqx_2 * best_d; + this_idx = idx2; + result = true; + } + return result; +} +float compute_1block_iq2ks_rmse(float d, const __m256 * vx, const __m256 * vw, const int8_t * values, __m256i & this_idx) { + constexpr int kBlockSize = 32; + uint32_t aux32; + std::memcpy(&aux32, values, sizeof(aux32)); + auto ivalues = _mm256_set1_epi32(aux32); + __m256 vbest[4]; + __m256i ibest[4]; + auto val = _mm256_set1_ps(d*values[0]); + auto ival = _mm256_set1_epi32(0); + for (int k = 0; k < kBlockSize/8; ++k) { + auto diff = _mm256_sub_ps(vx[k], val); + vbest[k] = _mm256_mul_ps(diff, diff); + ibest[k] = ival; + } + for (int j = 1; j < 4; ++j) { + val = _mm256_set1_ps(d*values[j]); + ival = _mm256_set1_epi32(j); + for (int k = 0; k < kBlockSize/8; ++k) { + auto diff = _mm256_sub_ps(vx[k], val); + diff = _mm256_mul_ps(diff, diff); + auto mask = _mm256_cmp_ps(diff, vbest[k], _CMP_LT_OQ); + vbest[k] = _mm256_or_ps(_mm256_and_ps(mask, diff), _mm256_andnot_ps(mask, vbest[k])); + auto imask = _mm256_castps_si256(mask); + ibest[k] = _mm256_or_si256(_mm256_and_si256(imask, ival), _mm256_andnot_si256(imask, ibest[k])); + } + } + auto idx = to_int8(ibest); + to_values_i32(idx, ivalues, ibest); + auto vd = _mm256_set1_ps(-d); + auto vrmse = _mm256_setzero_ps(); + for (int k = 0; k < 4; ++k) { + auto vq = _mm256_cvtepi32_ps(ibest[k]); + auto diff = _mm256_fmadd_ps(vd, vq, vx[k]); + auto wdiff = _mm256_mul_ps(vw[k], diff); + vrmse = _mm256_fmadd_ps(wdiff, diff, vrmse); + } + this_idx = idx; + return hsum_float_8(vrmse); +} +void quantize_row_iq2_ks_fast_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, float * all_scales, float * all_sw, int8_t * all_Ls) { + + constexpr int kBlockSize = 32; + + ggml_half * dptr = (ggml_half *)vy; + *dptr = GGML_FP32_TO_FP16(0.f); + + block_iq2_ks * y = (block_iq2_ks *)(dptr + 1); + + float weight[kBlockSize]; + + const int8_t * shifted_values = iq2nl_values + 4; + + const int nblock = n_per_row/QK_K; + + __m256 vx[4], vw[4]; + + for (int ibl = 0; ibl < nblock; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq2_ks)); + + auto scales = all_scales + ibl*(QK_K/kBlockSize); + auto sw = all_sw + ibl*(QK_K/kBlockSize); + + const float * xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 1.5f*sumx2/QK_K; + + uint16_t extra = 0; + + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + const float * xb = xbl + kBlockSize*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float amax = 0, max = 0, sumw = 0; + for (int j = 0; j < kBlockSize; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + sumw += weight[j]; + } + sw[ib] = sumw; + if (amax < 1e-14f) { + scales[ib] = 0; + continue; + } + for (int k = 0; k < 4; ++k) { + vx[k] = _mm256_loadu_ps(xb + 8*k); + vw[k] = _mm256_loadu_ps(weight + 8*k); + } + float d = max/iq2nl_values[7]; + float best = 0; + __m256i this_idx; + float this_d, this_score; + if (compute_1block_iq2ks(d, vx, vw, iq2nl_values, this_idx, this_d, this_score)) { + best = this_score; d = this_d; + } + for (int itry = -13; itry <= 13; ++itry) { + if (compute_1block_iq2ks(max/(iq2nl_values[0] + 0.5f*itry), vx, vw, iq2nl_values, this_idx, this_d, this_score)) { + if (this_score > best) { + best = this_score; d = this_d; + } + } + } + bool is_shifted = false; + for (int itry = -13; itry <= 13; ++itry) { + if (compute_1block_iq2ks(max/(iq2nl_values[4] + 0.5f*itry), vx, vw, iq2nl_values + 4, this_idx, this_d, this_score)) { + if (this_score > best) { + best = this_score; d = this_d; is_shifted = true; + } + } + } + scales[ib] = d; + if (is_shifted) extra |= (1 << ib); + } + y[ibl].extra = extra; + } + + float d = make_qx_quants(nblock*(QK_K/kBlockSize), 16, all_scales, all_Ls, all_sw); + + if (!d) return; + + auto vsumqx = _mm256_setzero_ps(); + auto vsumq2 = _mm256_setzero_ps(); + for (int ibl = 0; ibl < nblock; ++ibl) { + auto scales = all_scales + ibl*(QK_K/kBlockSize); + auto xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 1.5f*sumx2/QK_K; + auto Ls = all_Ls + ibl*(QK_K/kBlockSize); + __m256i idx[4]; + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + const int8_t * block_values = y[ibl].extra & (1 << ib) ? shifted_values : iq2nl_values; + uint32_t aux32; + std::memcpy(&aux32, block_values, sizeof(aux32)); + auto ivalues = _mm256_set1_epi32(aux32); + const float * xb = xbl + kBlockSize*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + for (int k = 0; k < 4; ++k) { + vx[k] = _mm256_loadu_ps(xb + 8*k); + vw[k] = _mm256_loadu_ps(weight + 8*k); + } + int ls = Ls[ib] - 16; + float dl = d*ls; + __m256i idx1, idx2; + auto rmse1 = compute_1block_iq2ks_rmse(dl, vx, vw, block_values, idx1); + if (Ls[ib] > 0 && dl > scales[ib]) { + auto rmse2 = compute_1block_iq2ks_rmse(d*(Ls[ib] - 17), vx, vw, block_values, idx2); + if (rmse2 < rmse1) { + --Ls[ib]; idx1 = idx2; + } + } + else if (Ls[ib] < 15 && dl < scales[ib]) { + auto rmse2 = compute_1block_iq2ks_rmse(d*(Ls[ib] - 15), vx, vw, block_values, idx2); + if (rmse2 < rmse1) { + ++Ls[ib]; idx1 = idx2; + } + } + __m256i iv[4]; + to_values_i32(idx1, ivalues, iv); + auto vd = _mm256_set1_ps(Ls[ib] - 16); + for (int k = 0; k < 4; ++k) { + auto vq = _mm256_mul_ps(vd, _mm256_cvtepi32_ps(iv[k])); + auto wvq = _mm256_mul_ps(vw[k], vq); + vsumqx = _mm256_fmadd_ps(wvq, vx[k], vsumqx); + vsumq2 = _mm256_fmadd_ps(wvq, vq, vsumq2); + } + ls = Ls[ib]; + y[ibl].scales[ib/2] |= ((ls & 0xf) << 4*(ib%2)); + y[ibl].extra |= ((ls >> 4) << (8 + ib)); + idx[ib % 4] = idx1; + if ((ib % 4) == 3) { + auto vqs1 = _mm256_or_si256(idx[0], _mm256_slli_epi16(idx[1], 2)); + auto vqs2 = _mm256_or_si256(_mm256_slli_epi16(idx[2], 4), _mm256_slli_epi16(idx[3], 6)); + auto vqs = _mm256_or_si256(vqs1, vqs2); + _mm256_storeu_si256((__m256i *)y[ibl].qs + ib/4, vqs); + } + } + } + float sumqx = hsum_float_8(vsumqx); + float sumq2 = hsum_float_8(vsumq2); + *dptr = GGML_FP32_TO_FP16(1.000f*(sumq2 > 0 ? sumqx/sumq2 : d)); +} +#endif +void quantize_row_iq2_ks_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, float * all_scales, float * all_sw, int8_t * all_Ls) { + + constexpr int kBlockSize = 32; + constexpr int kMax_i1 = 3*kBlockSize/4; + constexpr int kMin_i3 = kBlockSize/4; + + ggml_half * dptr = (ggml_half *)vy; + *dptr = GGML_FP32_TO_FP16(0.f); + + block_iq2_ks * y = (block_iq2_ks *)(dptr + 1); + + float weight[kBlockSize]; + float sumx[kBlockSize+1], sumw[kBlockSize+1]; + + std::array, kBlockSize> pairs; + + float val [4] = {float(iq2nl_values[0]), float(iq2nl_values[1]), float(iq2nl_values[2]), float(iq2nl_values[3])}; + float sval[4] = {float(iq2nl_values[4]), float(iq2nl_values[5]), float(iq2nl_values[6]), float(iq2nl_values[7])}; + + const int8_t * shifted_values = iq2nl_values + 4; + + const int nblock = n_per_row/QK_K; + + for (int ibl = 0; ibl < nblock; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq2_ks)); + + auto scales = all_scales + ibl*(QK_K/kBlockSize); + auto sw = all_sw + ibl*(QK_K/kBlockSize); + + const float * xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 1.5f*sumx2/QK_K; + + uint16_t extra = 0; + + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + const float * xb = xbl + kBlockSize*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + sw[ib] = 0; + float amax = 0; + for (int j = 0; j < kBlockSize; ++j) { + sw[ib] += weight[j]; + pairs[j] = {xb[j], j}; + float ax = std::abs(xb[j]); + amax = std::max(amax, ax); + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + std::sort(pairs.begin(), pairs.end()); + sumx[0] = sumw[0] = 0; + for (int j = 0; j < kBlockSize; ++j) { + int jj = pairs[j].second; + sumw[j+1] = sumw[j] + weight[jj]; + sumx[j+1] = sumx[j] + weight[jj]*xb[jj]; + } + float best = 0, d = 0; + bool is_shifted = false; + float sumqx, sumq2; + for (int i1 = 0; i1 < kMax_i1; ++i1) { + for (int i2 = i1; i2 < kBlockSize; ++i2) { + for (int i3 = std::max(i2, kMin_i3); i3 < kBlockSize; ++i3) { + sumqx = (sumx[i1] - sumx[ 0])*val[0] + (sumx[i2] - sumx[i1])*val[1] + + (sumx[i3] - sumx[i2])*val[2] + (sumx[kBlockSize] - sumx[i3])*val[3]; + sumq2 = (sumw[i1] - sumw[ 0])*val[0]*val[0] + (sumw[i2] - sumw[i1])*val[1]*val[1] + + (sumw[i3] - sumw[i2])*val[2]*val[2] + (sumw[kBlockSize] - sumw[i3])*val[3]*val[3]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; is_shifted = false; + } + sumqx = (sumx[i1] - sumx[ 0])*sval[0] + (sumx[i2] - sumx[i1])*sval[1] + + (sumx[i3] - sumx[i2])*sval[2] + (sumx[kBlockSize] - sumx[i3])*sval[3]; + sumq2 = (sumw[i1] - sumw[ 0])*sval[0]*sval[0] + (sumw[i2] - sumw[i1])*sval[1]*sval[1] + + (sumw[i3] - sumw[i2])*sval[2]*sval[2] + (sumw[kBlockSize] - sumw[i3])*sval[3]*sval[3]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; is_shifted = true; + } + sumqx = (sumx[i1] - sumx[ 0])*val[3] + (sumx[i2 ] - sumx[i1])*val[2] + + (sumx[i3] - sumx[i2])*val[1] + (sumx[kBlockSize] - sumx[i3])*val[0]; + sumq2 = (sumw[i1] - sumw[ 0])*val[3]*val[3] + (sumw[i2 ] - sumw[i1])*val[2]*val[2] + + (sumw[i3] - sumw[i2])*val[1]*val[1] + (sumw[kBlockSize] - sumw[i3])*val[0]*val[0]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; is_shifted = false; + } + sumqx = (sumx[i1] - sumx[ 0])*sval[3] + (sumx[i2 ] - sumx[i1])*sval[2] + + (sumx[i3] - sumx[i2])*sval[1] + (sumx[kBlockSize] - sumx[i3])*sval[0]; + sumq2 = (sumw[i1] - sumw[ 0])*sval[3]*sval[3] + (sumw[i2 ] - sumw[i1])*sval[2]*sval[2] + + (sumw[i3] - sumw[i2])*sval[1]*sval[1] + (sumw[kBlockSize] - sumw[i3])*sval[0]*sval[0]; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; is_shifted = true; + } + } + } + } + scales[ib] = d; + if (is_shifted) extra |= (1 << ib); + } + y[ibl].extra = extra; + } + + float d = make_qx_quants(nblock*(QK_K/kBlockSize), 16, all_scales, all_Ls, all_sw); + + if (!d) return; + + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 1.5f*sumx2/QK_K; + auto Ls = all_Ls + ibl*(QK_K/kBlockSize); + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + int ls = Ls[ib]; + y[ibl].scales[ib/2] |= ((ls & 0xf) << 4*(ib%2)); + y[ibl].extra |= ((ls >> 4) << (8 + ib)); + ls -= 16; + float dl = d * ls; + if (dl) { + const int8_t * block_values = y[ibl].extra & (1 << ib) ? shifted_values : iq2nl_values; + const float * xb = xbl + kBlockSize*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float idl = 1/dl; + uint8_t * qs = y[ibl].qs + 32*(ib/4); + for (int j = 0; j < 32; ++j) { + const float al = idl*xb[j]; + int ibest = best_index_iq2nl(block_values, al); + qs[j] |= (ibest << 2*(ib%4)); + float w = weight[j]; + float q = block_values[ibest]*ls; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + } + } + } + *dptr = GGML_FP32_TO_FP16(1.030f*(sumq2 > 0 ? sumqx/sumq2 : d)); +} +} + +void quantize_row_iq2_ks_ref(const float * x, block_iq2_ks * y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq2_ks(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq2_ks(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + block_iq2_ks * y = (block_iq2_ks *)vy; + quantize_row_iq2_ks_ref(x, y, k); +} + +size_t quantize_iq2_ks(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + constexpr int kBlockSize = 32; + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ2_KS, n_per_row); + int nblock = n_per_row/QK_K; + std::vector all_scales(nblock*(QK_K/kBlockSize)), all_sw(nblock*(QK_K/kBlockSize)); + std::vector all_Ls(nblock*(QK_K/kBlockSize)); + auto q_func = [&all_scales, &all_sw, &all_Ls] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { +#ifdef __AVX2__ + if (user_data && !user_data->slow_iq2_ks) { + quantize_row_iq2_ks_fast_impl(x, vy, n_per_row, imatrix, all_scales.data(), all_sw.data(), all_Ls.data()); + return; + } +#endif + quantize_row_iq2_ks_impl(x, vy, n_per_row, imatrix, all_scales.data(), all_sw.data(), all_Ls.data()); + }; + QHelper helper(imatrix, user_data, n_per_row, kBlockSize); + helper.quantize(nrows, src, dst, row_size, q_func); + return nrows * row_size; +} + +void dequantize_row_iq2_ks(const block_iq2_ks * x, float * y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + const ggml_half * dptr = (const ggml_half *)x; + const float d = GGML_FP16_TO_FP32(*dptr); + x = (const block_iq2_ks *)(dptr + 1); + + for (int i = 0; i < nb; i++) { + + const uint8_t * qs = x[i].qs; + + uint16_t extra = x[i].extra; + + int shift = 0; + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + float dl1 = d * (((x[i].scales[ib64] & 0xf) | ((extra >> 4) & 0x10)) - 16); + float dl2 = d * (((x[i].scales[ib64] >> 4) | ((extra >> 5) & 0x10)) - 16); + const int8_t * values1 = extra & 1 ? iq2nl_values + 4 : iq2nl_values; + const int8_t * values2 = extra & 2 ? iq2nl_values + 4 : iq2nl_values; + extra >>= 2; + for (int j = 0; j < 32; ++j) { + y[j+ 0] = dl1 * values1[(qs[j] >> (shift+0)) & 3]; + y[j+32] = dl2 * values2[(qs[j] >> (shift+2)) & 3]; + } + y += 64; + shift += 4; + if (shift == 8) { qs += 32; shift = 0; } + } + + } + +} + +void vec_dot_iq2_ks_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_KS, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + + const ggml_half * dptr = (const ggml_half *)vx; + const float d = GGML_FP16_TO_FP32(*dptr); + const block_iq2_ks * x = (const block_iq2_ks *)(dptr + 1); + const block_q8_K * y = (const block_q8_K *)vy; + + const int nb = n / QK_K; + float sumf = 0; + for (int i = 0; i < nb; i++) { + const uint8_t * qs = x[i].qs; + const int8_t * q8 = y[i].qs; + uint16_t extra = x[i].extra; + int sumi = 0; + for (int ib128 = 0; ib128 < QK_K/128; ++ib128) { + int d1 = (((x[i].scales[2*ib128+0] & 0xf) | ((extra >> 4) & 0x10)) - 16); + int d2 = (((x[i].scales[2*ib128+0] >> 4) | ((extra >> 5) & 0x10)) - 16); + int d3 = (((x[i].scales[2*ib128+1] & 0xf) | ((extra >> 6) & 0x10)) - 16); + int d4 = (((x[i].scales[2*ib128+1] >> 4) | ((extra >> 7) & 0x10)) - 16); + const int8_t * values1 = extra & 1 ? iq2nl_values + 4 : iq2nl_values; + const int8_t * values2 = extra & 2 ? iq2nl_values + 4 : iq2nl_values; + const int8_t * values3 = extra & 4 ? iq2nl_values + 4 : iq2nl_values; + const int8_t * values4 = extra & 8 ? iq2nl_values + 4 : iq2nl_values; + extra >>= 4; + int sumi1 = 0, sumi2 = 0, sumi3 = 0, sumi4 = 0; + for (int j = 0; j < 32; ++j) { + sumi1 += q8[j+ 0] * values1[(qs[j] >> 0) & 3]; + sumi2 += q8[j+32] * values2[(qs[j] >> 2) & 3]; + sumi3 += q8[j+64] * values3[(qs[j] >> 4) & 3]; + sumi4 += q8[j+96] * values4[(qs[j] >> 6) & 3]; + } + sumi += d1*sumi1 + d2*sumi2 + d3*sumi3 + d4*sumi4; + q8 += 128; + qs += 32; + } + sumf += y[i].d * sumi; + } + + *s = d * sumf; + +} + +// +// ======================================== iq2_kl +// +namespace { + +const int8_t iq3nl_index[111] = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 8, 8, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 9, + 9, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 10, 10, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 11, 11, 4, 4, 4, 4, + 4, 4, 4, 4, 4, 4, 12, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 13, 13, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, + 6, 6, 6, 6, 14, 14, 7, 7, 7, 7, 7, 7, 7, 7, 7 +}; +inline int best_index_iq3nl(const int8_t * values, float x) { + int ix = (int)x - values[0]; + if (ix < 0 || ix >= 111) return ix < 0 ? 0 : 7; + ix = iq3nl_index[ix]; + return ix < 8 ? ix : x - values[ix-8] < values[ix-7] - x ? ix-8 : ix-7; +} + +void quantize_row_iq2_kl_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, float * all_scales) { + constexpr int kBlockSize = 32; + constexpr float kSigmaFactor = 2.25f; + constexpr int ntry = 5; + static const int k_index[64] = {-1, -2, 0, -3, -4, 1, -5, -6, 2, -7, -8, 3, -9, 4, -10, 5, -11, 6, 7, -12, 8, 9, 10, -13, 11, -14, -15, -16, 12, 13, -17, + 14, -18, -19, 15, 16, 17, 18, 19, -20, -21, 20, 21, 22, 23, 24, -22, -23, 25, -24, 26, -25, 27, -26, 28, 29, -27, -28, 30, -29, -30, 31, -31, -32}; + static const std::vector> k_neighbours = { + { 2, 0, 6, 11, 7, 3, 8, 15, }, + { 0, 2, 3, 6, 7, 1, 8, 4, }, + { 0, 1, 3, 4, 8, 7, 9, 6, }, + { 1, 0, 3, 4, 8, 9, 7, 10, }, + { 1, 4, 5, 10, 9, 3, 8, 0, }, + { 5, 1, 4, 10, 9, 14, 8, 3, }, + { 6, 2, 7, 0, 3, 11, 8, 15, }, + { 3, 7, 0, 6, 8, 4, 12, 9, }, + { 3, 4, 8, 9, 1, 7, 12, 10, }, + { 4, 10, 5, 9, 1, 8, 13, 14, }, + { 11, 2, 6, 7, 20, 15, 25, 21, }, + { 8, 7, 3, 12, 9, 16, 17, 13, }, + { 14, 5, 10, 19, 9, 13, 4, 18, }, + { 6, 15, 7, 11, 20, 21, 16, 2, }, + { 15, 7, 16, 6, 21, 12, 17, 22, }, + { 12, 16, 17, 8, 15, 7, 13, 22, }, + { 19, 10, 13, 18, 14, 9, 12, 24, }, + { 11, 20, 25, 6, 15, 2, 21, 7, }, + { 20, 15, 21, 6, 11, 7, 16, 26, }, + { 14, 19, 29, 10, 28, 18, 13, 24, }, + { 25, 11, 20, 21, 15, 6, 26, 30, }, + { 19, 24, 28, 18, 29, 23, 13, 17, }, + { 29, 19, 14, 28, 24, 18, 10, 13, }, + { 20, 26, 21, 25, 30, 15, 22, 16, }, + { 27, 26, 22, 23, 21, 30, 16, 24, }, + { 27, 24, 28, 31, 23, 18, 22, 17, }, + { 25, 30, 20, 26, 21, 11, 15, 22, }, + { 30, 26, 25, 20, 21, 27, 22, 15, }, + { 30, 27, 31, 26, 22, 23, 21, 24, }, + { 31, 27, 30, 26, 28, 23, 22, 24, }, + { 31, 28, 29, 27, 24, 23, 19, 18, }, + { 29, 28, 31, 24, 19, 27, 14, 18, }, + }; + auto values = iq3nl_values; + std::pair grid[32]; + for (int j = 0; j < 64; ++j) { + if (int i = k_index[j]; i >= 0) { + int i1 = j/8, i2 = j%8; + grid[i] = {values[i1], values[i2]}; + } + } + + ggml_half * dptr = (ggml_half *)vy; + auto y = (block_iq2_kl *)(dptr + 1); + + float weight[kBlockSize]; + + auto index = [&grid, values] (float id, float x1, float x2, float w1, float w2) { + float sx1 = id*x1; + float sx2 = id*x2; + int l1 = best_index_iq3nl(values, sx1); + int l2 = best_index_iq3nl(values, sx2); + int i = k_index[8*l1 + l2]; + if (i >= 0) return i; + auto& neigh = k_neighbours[-i-1]; + float best = std::numeric_limits::max(); + int ibest = -1; + for (auto& n : neigh) { + float diff1 = grid[n].first - sx1; + float diff2 = grid[n].second - sx2; + float score = w1*diff1*diff1 + w2*diff2*diff2; + if (score < best) { + best = score; ibest = n; + } + } + GGML_ASSERT(ibest >= 0); + return ibest; + }; + + float max_scale = 0, max_abs_scale = 0; + + for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { + std::memset(&y[ibl], 0, sizeof(block_iq2_kl)); + auto scales = all_scales + ibl*(QK_K/kBlockSize); + auto xbl = x + ibl*QK_K; + float sigma2 = 0; + for (int j = 0; j < QK_K; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= kSigmaFactor/QK_K; + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + auto xb = xbl + ib*kBlockSize; + if (quant_weights) { + auto qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j]*sqrt(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = std::abs(xb[j]); //xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < kBlockSize; ++j) { + float ax = std::abs(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/values[0] : max/values[0]; + float id = 1/d; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + for (int j = 0; j < kBlockSize; j += 2) { + float w1 = weight[j+0]; + float w2 = weight[j+1]; + int idx = index(id, xb[j+0], xb[j+1], w1, w2); + float q1 = grid[idx].first ; + float q2 = grid[idx].second; + sumqx_p += w1*q1*xb[j] + w2*q2*xb[j+1]; + sumq2_p += w1*q1*q1 + w2*q2*q2; + idx = index(-id, xb[j+0], xb[j+1], w1, w2); + q1 = grid[idx].first ; + q2 = grid[idx].second; + sumqx_m += w1*q1*xb[j] + w2*q2*xb[j+1]; + sumq2_m += w1*q1*q1 + w2*q2*q2; + } + d = sumqx_p/sumq2_p; + float best = d*sumqx_p; + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d*sumqx_m; + } + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry + values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < kBlockSize; j += 2) { + float w1 = weight[j+0]; + float w2 = weight[j+1]; + int idx = index(id, xb[j+0], xb[j+1], w1, w2); + float q1 = grid[idx].first ; + float q2 = grid[idx].second; + sumqx_p += w1*q1*xb[j] + w2*q2*xb[j+1]; + sumq2_p += w1*q1*q1 + w2*q2*q2; + idx = index(-id, xb[j+0], xb[j+1], w1, w2); + q1 = grid[idx].first ; + q2 = grid[idx].second; + sumqx_m += w1*q1*xb[j] + w2*q2*xb[j+1]; + sumq2_m += w1*q1*q1 + w2*q2*q2; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; + } + } + scales[ib] = d; + float ad = std::abs(d); + if (ad > max_abs_scale) { + max_abs_scale = ad; max_scale = d; + } + } + } + + if (!max_abs_scale) { + dptr[0] = GGML_FP32_TO_FP16(0.f); + return; + } + + float d = -max_scale/32; + float id = 1/d; + + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { + auto scales = all_scales + ibl*(QK_K/kBlockSize); + auto xbl = x + ibl*QK_K; + float sigma2 = 0; + for (int j = 0; j < QK_K; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= kSigmaFactor/QK_K; + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + auto xb = xbl + ib*kBlockSize; + if (quant_weights) { + auto qw = quant_weights + ibl*QK_K + ib*kBlockSize; + for (int j = 0; j < kBlockSize; ++j) weight[j] = qw[j]*sqrt(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = std::abs(xb[j]); //xb[j]*xb[j]; + } + int ls = nearest_int(id*scales[ib]); + ls = std::max(-32, std::min(31, ls)); + int lsmin = std::max(-32, ls-1); + int lsmax = std::min( 31, ls+1); + float best_score = std::numeric_limits::max(); + int best_ls = ls; + for (int ils = lsmin; ils <= lsmax; ++ils) { + float dl = d*ils; + float idl = dl ? 1/dl : 0.f; + float score = 0; + for (int j = 0; j < kBlockSize/2; ++j) { + float w1 = weight[2*j+0]; + float w2 = weight[2*j+1]; + int idx = index(idl, xb[2*j+0], xb[2*j+1], w1, w2); + float diff1 = dl*grid[idx].first - xb[2*j+0]; + float diff2 = dl*grid[idx].second - xb[2*j+1]; + score += w1*diff1*diff1 + w2*diff2*diff2; + } + if (score < best_score) { + best_score = score; + best_ls = ils; + } + } + ls = best_ls; + int uls = ls + 32; + y[ibl].scales_l[ib%4] |= ((uls & 0xf) << 4*(ib/4)); + y[ibl].scales_h |= ((uls >> 4) << 2*ib); + if (ls == 0) continue; + float dl = d*ls; + float idl = 1/dl; + for (int j = 0; j < kBlockSize/2; ++j) { + float w1 = weight[2*j+0]; + float w2 = weight[2*j+1]; + int idx = index(idl, xb[2*j+0], xb[2*j+1], w1, w2); + y[ibl].qs[16*(ib/2) + j] |= ((idx & 0xf) << 4*(ib%2)); + y[ibl].qh[j] |= ((idx >> 4) << ib); + float q1 = ls*grid[idx].first ; + float q2 = ls*grid[idx].second; + sumqx += w1*q1*xb[2*j] + w2*q2*xb[2*j+1]; + sumq2 += w1*q1*q1 + w2*q2*q2; + } + } + } + if (sumq2 > 0) d = sumqx/sumq2; + + dptr[0] = GGML_FP32_TO_FP16(1.025f * d); + +} +} + +void quantize_row_iq2_kl_ref(const float * x, block_iq2_kl * y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq2_kl(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq2_kl(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + block_iq2_kl * y = (block_iq2_kl *)vy; + quantize_row_iq2_kl_ref(x, y, k); +} + +size_t quantize_iq2_kl(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + constexpr int kBlockSize = 32; + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ2_KL, n_per_row); + int nblock = n_per_row/QK_K; + std::vector all_scales(nblock*(QK_K/kBlockSize)); + auto q_func = [&all_scales] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + quantize_row_iq2_kl_impl(x, vy, n_per_row, imatrix, all_scales.data()); + }; + QHelper helper(imatrix, user_data, n_per_row, kBlockSize); + helper.quantize(nrows, src, dst, row_size, q_func); + return nrows * row_size; +} + +void dequantize_row_iq2_kl(const block_iq2_kl * x, float * y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + const ggml_half * dptr = (const ggml_half *)x; + const float d = GGML_FP16_TO_FP32(*dptr); + x = (const block_iq2_kl *)(dptr + 1); + + for (int i = 0; i < nb; i++) { + + auto qs = x[i].qs; + auto qh = x[i].qh; + auto scales_h = x[i].scales_h; + + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + float dl1 = d * (int(((x[i].scales_l[(2*ib64+0)%4] >> 4*(ib64/2)) & 0xf) | (((scales_h >> (4*ib64+0)) & 3) << 4)) - 32); + float dl2 = d * (int(((x[i].scales_l[(2*ib64+1)%4] >> 4*(ib64/2)) & 0xf) | (((scales_h >> (4*ib64+2)) & 3) << 4)) - 32); + for (int j = 0; j < 16; ++j) { + const int8_t * val1 = (const int8_t *)(iq2kl_values + ((qs[j] & 0xf) | (((qh[j] >> (2*ib64+0)) & 1) << 4))); + const int8_t * val2 = (const int8_t *)(iq2kl_values + ((qs[j] >> 4) | (((qh[j] >> (2*ib64+1)) & 1) << 4))); + y[2*j+ 0] = dl1 * val1[0]; + y[2*j+ 1] = dl1 * val1[1]; + y[2*j+32] = dl2 * val2[0]; + y[2*j+33] = dl2 * val2[1]; + } + y += 64; + qs += 16; + } + + } +} + +void vec_dot_iq2_kl_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_KL, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif +} + +// +// ============================================== iq3_k +// +namespace { + +static void quantize_row_iq3_k_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, + [[maybe_unused]] const quantize_user_data * user_data) { + + constexpr int ntry = 3; + + block_iq3_k * y = (block_iq3_k *)vy; + + float scales[QK_K/16]; + float weight[16]; + uint8_t L[16]; + + const int8_t * shifted_values = iq3nl_values + 8; + + for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq3_k)); + y[ibl].d = GGML_FP32_TO_FP16(0.f); + + const float * xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 1.5f*sumx2/QK_K; + + uint16_t extra = 0; + + float max_abs_scale = 0; + + for (int ib = 0; ib < QK_K/16; ++ib) { + const float * xb = xbl + 16*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*16; + for (int j = 0; j < 16; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < 16; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < 16; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/iq3nl_values[0] : max/iq3nl_values[0]; + float id = 1/d; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + float best = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq3nl(iq3nl_values, al); + float q = iq3nl_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq3nl(iq3nl_values, -al); + q = iq3nl_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0) { + d = sumqx_p/sumq2_p; + best = d*sumqx_p; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d*sumqx_m; + } + bool is_shifted = false; + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (2*itry + iq3nl_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq3nl(iq3nl_values, al); + float q = iq3nl_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq3nl(iq3nl_values, -al); + q = iq3nl_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = false; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = false; + } + id = (2*itry + shifted_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq3nl(shifted_values, al); + float q = shifted_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq3nl(shifted_values, -al); + q = shifted_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = true; + } + } + if (!d) { + scales[ib] = 0; continue; + } + + const int8_t * block_values = is_shifted ? shifted_values : iq3nl_values; + float sumqx = 0, sumq2 = 0; + id = 1/d; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq3nl(block_values, al); + L[j] = l; + float q = block_values[l]; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + if (sumq2 > 0) d = sumqx/sumq2; + + float best_d = d; + for (int iter = 0; iter < 128; ++iter) { + float gmax = 0; + int best_j = -1, dir = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float g = d * w * (xb[j] - d*block_values[L[j]]); + if (g > 0 && L[j] < 7) { + if (g > gmax) { + gmax = g; best_j = j; dir = 1; + } + } + else if (g < 0 && L[j] > 0) { + if (-g > gmax) { + gmax = -g; best_j = j; dir = -1; + } + } + } + if (best_j < 0) break; + + float w = weight[best_j]; + sumqx += w*xb[best_j]*(block_values[L[best_j]+dir] - block_values[L[best_j]]); + sumq2 += w*(block_values[L[best_j]+dir]*block_values[L[best_j]+dir] - block_values[L[best_j]]*block_values[L[best_j]]); + L[best_j] += dir; + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + best_d = sumqx/sumq2; best = best_d*sumqx; + } + else if (iter > 8) break; + + } + + scales[ib] = d; + + if (is_shifted) extra |= (1 << ib); + + float abs_scale = fabsf(scales[ib]); + max_abs_scale = MAX(max_abs_scale, abs_scale); + } + + if (!max_abs_scale) continue; + + float d = max_abs_scale/31; + y[ibl].extra = extra; + float id = 1/d; + + float sumqx = 0, sumq2 = 0; + for (int ib = 0; ib < QK_K/16; ++ib) { + int ls = nearest_int(0.5f*(id*fabsf(scales[ib])-1)); + ls = MAX(0, MIN(15, ls)); + y[ibl].scales_l[ib/2] |= (ls << 4*(ib%2)); + if (scales[ib] < 0) y[ibl].scales_h |= (1 << ib); + ls = (2*ls + 1) * (scales[ib] < 0 ? -1 : 1); + float dl = d * ls; + if (dl) { + const int8_t * block_values = y[ibl].extra & (1 << ib) ? shifted_values : iq3nl_values; + const float * xb = xbl + 16*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*16; + for (int j = 0; j < 16; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < 16; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float idl = 1/dl; + int ib32 = ib/2; + int offset = 16*(ib%2); + uint8_t * qs = y[ibl].qs + 32*(ib32/4) + offset; + uint8_t * qh = y[ibl].qh + 32*(ib32/8) + offset; + for (int j = 0; j < 16; ++j) { + const float al = idl*xb[j]; + int ibest = best_index_iq3nl(block_values, al); + qs[j] |= ((ibest & 3) << 2*(ib32%4)); + qh[j] |= ((ibest >> 2) << (ib32%8)); + float w = weight[j]; + float q = block_values[ibest]*ls; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + } + } + y[ibl].d = GGML_FP32_TO_FP16(1.01f*(sumq2 > 0 ? sumqx/sumq2 : d)); + + } +} + +} + +void quantize_row_iq3_k_ref(const float * x, block_iq3_k * y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq3_k(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq3_k(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + block_iq3_k * y = (block_iq3_k *)vy; + quantize_row_iq3_k_ref(x, y, k); +} + +size_t quantize_iq3_k(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + QHelper helper(imatrix, user_data, n_per_row, 16); + auto row_size = ggml_row_size(GGML_TYPE_IQ3_K, n_per_row); + helper.quantize(nrows, src, dst, row_size, quantize_row_iq3_k_impl); + return nrows * row_size; +} + +void dequantize_row_iq3_k(const block_iq3_k * x, float * y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + + uint16_t sh = x[i].scales_h; + uint16_t extra = x[i].extra; + + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + float dl1 = d * ((2*(x[i].scales_l[ib32] & 0xf) + 1) * ((sh & 1) ? -1 : 1)); + float dl2 = d * ((2*(x[i].scales_l[ib32] >> 4) + 1) * ((sh & 2) ? -1 : 1)); + sh >>= 2; + const int8_t * values1 = extra & 1 ? iq3nl_values + 8 : iq3nl_values; + const int8_t * values2 = extra & 2 ? iq3nl_values + 8 : iq3nl_values; + extra >>= 2; + int shift_l = 2*(ib32%4); + int shift_h = ib32%8; + for (int j = 0; j < 16; ++j) { + y[j+ 0] = dl1 * values1[((qs[j+ 0] >> shift_l) & 3) | (((qh[j+ 0] >> shift_h) & 1) << 2)]; + y[j+16] = dl2 * values2[((qs[j+16] >> shift_l) & 3) | (((qh[j+16] >> shift_h) & 1) << 2)]; + } + y += 32; + if (shift_l == 6) qs += 32; + } + + } +} + +void vec_dot_iq3_k_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ3_K, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + + GGML_ABORT("not implemented"); +} + +// +// ============================================== iq3_ks +// +namespace { +static void quantize_row_iq3_ks_impl(const int super_block_size, const int block_size, + int n_per_row, const float * x, char * cy, + float * all_scales, float * weight, + const int8_t * values, + const float * quant_weights, + const int ntry) { + + ggml_half * dptr = (ggml_half *)cy; + block_iq3_ks * y = (block_iq3_ks *)(dptr + 1); + + const int8_t * shifted_values = values + 8; + + float amax_scale = 0; + float max_scale = 0; + + for (int ibl = 0; ibl < n_per_row/super_block_size; ++ibl) { + memset(&y[ibl], 0, sizeof(block_iq3_ks)); + const float * xbl = x + ibl*super_block_size; + auto scales = all_scales + ibl*(super_block_size/block_size); + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= 2.f/super_block_size; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const float * xb = xbl + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ibl*super_block_size + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < block_size; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/values[0] : max/values[0]; + float id = 1/d; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + float best = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq3nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq3nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0) { + d = sumqx_p/sumq2_p; + best = d*sumqx_p; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d*sumqx_m; + } + bool is_shifted = false; + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry + values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq3nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq3nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = false; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = false; + } + id = (itry + shifted_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq3nl(shifted_values, al); + float q = shifted_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq3nl(shifted_values, -al); + q = shifted_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = true; + } + } + if (is_shifted) y[ibl].extra |= (1 << (8 + ib)); + scales[ib] = d; + float ascale = std::abs(d); + if (ascale > amax_scale) { + amax_scale = ascale; max_scale = d; + } + } + } + float d = -max_scale/16; + *dptr = GGML_FP32_TO_FP16(d); + if (!d) return; + float id = d ? 1/d : 0.f; + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < n_per_row/super_block_size; ++ibl) { + const float * xbl = x + ibl*super_block_size; + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= 2.f/super_block_size; + auto scales = all_scales + (super_block_size/block_size)*ibl; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const int8_t * block_values = (y[ibl].extra >> (8 + ib)) & 0x01 ? shifted_values : values; + int l = nearest_int(id*scales[ib]); + l = std::max(-16, std::min(15, l)); + uint8_t ul = l + 16; + y[ibl].scales[ib%4] |= (ul & 0xf) << 4*(ib/4); + y[ibl].extra |= (ul >> 4) << ib; + float dl = d * l; + float idl = dl ? 1/dl : 0.f; + const float * xb = xbl + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ibl*super_block_size + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + auto qs = y[ibl].qs + (ib/4)*block_size; + auto qh = y[ibl].qh + (ib/8)*block_size; + for (int j = 0; j < block_size; ++j) { + uint8_t i = best_index_iq3nl(block_values, idl*xb[j]); + qs[j] |= ((i & 3) << 2*(ib%4)); + qh[j] |= ((i >> 2) << (ib%8)); + float w = weight[j]; + float q = block_values[i]*l; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + } + } + if (sumq2 > 0) *dptr = GGML_FP32_TO_FP16(sumqx/sumq2); +} +} + +void quantize_row_iq3_ks_ref(const float * x, block_iq3_ks * y, int64_t k) { + quantize_iq3_ks(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq3_ks(const float * x, void * y, int64_t k) { + quantize_iq3_ks(x, (void *)y, 1, k, nullptr, nullptr); +} + +size_t quantize_iq3_ks(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + constexpr int kBlockSize = 32; + GGML_ASSERT(n_per_row%QK_K == 0); + float weight[kBlockSize]; + std::vector all_scales(n_per_row/kBlockSize); + auto row_size = ggml_row_size(GGML_TYPE_IQ3_KS, n_per_row); + QHelper helper(imatrix, user_data, n_per_row, kBlockSize); + auto q_func = [&all_scales, &weight, block_size = kBlockSize] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + quantize_row_iq3_ks_impl(QK_K, block_size, n_per_row, x, (char *)vy, all_scales.data(), weight, iq3nl_values, imatrix, 5); + }; + helper.quantize(nrows, src, dst, row_size, q_func); + return nrows * row_size; +} + +void dequantize_row_iq3_ks(const block_iq3_ks * x, float * y, int64_t k) { + constexpr int kBlockSize = 32; + static_assert(QK_K/kBlockSize == 8); + GGML_ASSERT(k%QK_K == 0); + const ggml_half * dptr = (const ggml_half *)x; + float d = GGML_FP16_TO_FP32(*dptr); + x = (const block_iq3_ks *)(dptr + 1); + float dl[8]; + int nblock = k/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int j = 0; j < 4; ++j) { + int ls1 = (x[ibl].scales[j] & 0xf) | (((x[ibl].extra >> (j+0)) & 1) << 4); + int ls2 = (x[ibl].scales[j] >> 4) | (((x[ibl].extra >> (j+4)) & 1) << 4); + dl[j+0] = d*(ls1 - 16); + dl[j+4] = d*(ls2 - 16); + } + auto qs = x[ibl].qs; + auto qh = x[ibl].qh; + for (int i128 = 0; i128 < QK_K/128; ++i128) { + for (int ib = 0; ib < 4; ++ib) { + const int8_t * values = iq3nl_values + ((x[ibl].extra >> (8 + (4*i128+ib)) & 1) << 3); + for (int j = 0; j < kBlockSize; ++j) { + y[j] = dl[4*i128 + ib] * values[((qs[j] >> 2*ib) & 3) | (((qh[j] >> (4*i128+ib)) & 1) << 2)]; + } + y += kBlockSize; + } + qs += kBlockSize; + } + } +} + +void vec_dot_iq3_ks_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ3_KS, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_ABORT("Not implemented"); +} + +// +// ============================================== iq4_K +// +void dequantize_row_iq4_k(const block_iq4_k * x, float * y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const uint8_t * qs = x[i].qs; + + const float d = GGML_FP16_TO_FP32(x[i].d); + + uint16_t extra = x[i].extra; + + for (int ib = 0; ib < QK_K/32; ++ib) { + const uint8_t sh = x[i].scales_h[ib/2] >> 4*(ib%2); + const float dl1 = d * (((x[i].scales_l[ib] & 0xf) | ((sh << 4) & 0x30)) - 32); + const float dl2 = d * (((x[i].scales_l[ib] >> 4) | ((sh << 2) & 0x30)) - 32); + const int8_t * values1 = extra & 1 ? iq4k_values + 16 : iq4k_values; + const int8_t * values2 = extra & 2 ? iq4k_values + 16 : iq4k_values; + extra >>= 2; + for (int j = 0; j < 16; ++j) { + y[j+ 0] = dl1 * values1[qs[j] & 0xf]; + y[j+16] = dl2 * values2[qs[j] >> 4]; + } + y += 32; + qs += 16; + } + } +} + +void vec_dot_iq4_k_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ4_K, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + + const int nb = n / QK_K; + + const block_iq4_k * x = (const block_iq4_k *)vx; + const block_q8_K * y = (const block_q8_K *)vy; + + float sumf = 0; + for (int ibl = 0; ibl < nb; ++ibl) { + const float d4d8 = GGML_FP16_TO_FP32(x[ibl].d) * y[ibl].d; + uint16_t extra = x[ibl].extra; + uint32_t h = *((const uint32_t *)x[ibl].scales_h); + const uint8_t * qs = x[ibl].qs; + const int8_t * q8 = y[ibl].qs; + int32_t sum = 0; + for (int ib = 0; ib < QK_K/32; ++ib) { + const int ls1 = ((x[ibl].scales_l[ib] & 0xf) | ((h << 4) & 0x30)) - 32; + const int ls2 = ((x[ibl].scales_l[ib] >> 4) | ((h << 2) & 0x30)) - 32; + h >>= 4; + const int8_t * values1 = iq4k_values + 16*(extra & 1); + const int8_t * values2 = iq4k_values + 8*(extra & 2); + extra >>= 2; + int sumi1 = 0, sumi2 = 0; + for (int j = 0; j < 16; ++j) { + sumi1 += q8[j+ 0] * values1[qs[j] & 0xf]; + sumi2 += q8[j+16] * values2[qs[j] >> 4]; + } + sum += ls1*sumi1 + ls2*sumi2; + qs += 16; + q8 += 32; + } + sumf += d4d8 * sum; + } + *s = sumf; + +} + +namespace { +const int8_t iq4nl_index[241] = { + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 16, 16, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 17, 17, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 18, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, + 3, 3, 3, 3, 3, 3, 19, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 20, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, + 5, 5, 21, 21, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 22, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 23, 23, 8, 8, 8, 8, + 8, 8, 8, 8, 8, 8, 24, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 25, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 26, 26, + 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 27, 27, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 28, 13, 13, 13, + 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 29, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, + 14, 14, 14, 14, 30, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15 +}; +inline int best_index_iq4nl(const int8_t * values, float x) { + int ix = (int)x - values[0]; + if (ix < 0 || ix >= 241) return ix < 0 ? 0 : 15; + ix = iq4nl_index[ix]; + return ix < 16 ? ix : x - values[ix-16] < values[ix-15] - x ? ix-16 : ix-15; +} + +static void quantize_row_iq4_k_impl_bs16(const int super_block_size, const int block_size, const float * x, + block_iq4_k * y, + float * scales, float * weight, uint8_t * L, + const int8_t * values, + const float * quant_weights, + const int ntry) { + + GGML_ASSERT(super_block_size == 256 && block_size == 16); + + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += x[j]*x[j]; + sigma2 *= 2.f/super_block_size; + + memset(y, 0, sizeof(block_iq4_k)); + y->d = GGML_FP32_TO_FP16(0.f); + + uint16_t * scales_h = (uint16_t *)y->scales_h; + + const int8_t * shifted_values = values + 16; + + float max_scale = 0, amax_scale = 0; + uint16_t extra = 0; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const float * xb = x + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < block_size; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/values[0] : max/values[0]; + float id = 1/d; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq4nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq4nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + d = sumqx_p/sumq2_p; + bool is_shifted = false; + float best = d*sumqx_p; + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d*sumqx_m; + } + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry + values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq4nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq4nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = false; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = false; + } + id = (itry + shifted_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq4nl(shifted_values, al); + float q = shifted_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq4nl(shifted_values, -al); + q = shifted_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = true; + } + } + if (is_shifted) extra |= (1 << ib); + scales[ib] = d; + float abs_d = fabsf(d); + if (abs_d > amax_scale) { + amax_scale = abs_d; max_scale = d; + } + } + float d = -max_scale/32; + y->d = GGML_FP32_TO_FP16(d); + y->extra = extra; + float id = d ? 1/d : 0.f; + float sumqx = 0, sumq2 = 0; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const int8_t * block_values = extra & (1 << ib) ? shifted_values : values; + int l = nearest_int(id*scales[ib]); + l = MAX(-32, MIN(31, l)); + float dl = d * l; + float idl = dl ? 1/dl : 0.f; + uint8_t * Lb = L + ib*block_size; + const float * xb = x + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + for (int j = 0; j < block_size; ++j) { + Lb[j] = best_index_iq4nl(block_values, idl*xb[j]); + float w = weight[j]; + float q = block_values[Lb[j]]*l; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + l += 32; + uint8_t l_l = l & 0xf; + uint8_t l_h = l >> 4; + if (ib%2 == 0) y->scales_l[ib/2] = l_l; + else y->scales_l[ib/2] |= (l_l << 4); + scales_h[ib/8] |= (l_h << 2*(ib%8)); + } + if (sumq2 > 0) y->d = GGML_FP32_TO_FP16(sumqx/sumq2); + + for (int i = 0; i < super_block_size/32; ++i) { + for (int j = 0; j < 16; ++j) { + y->qs[16*i + j] = L[32*i + j] | (L[32*i + 16 + j] << 4); + } + } +} + +} + +void quantize_row_iq4_k_ref(const float * x, block_iq4_k * y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq4_k(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq4_k(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + block_iq4_k * y = (block_iq4_k *)vy; + quantize_row_iq4_k_ref(x, y, k); +} + +size_t quantize_iq4_k(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + uint8_t L[QK_K]; + float weight[16]; + float scales[QK_K/16]; + auto q_func = [&L, &weight, &scales] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + block_iq4_k * iq4 = (block_iq4_k *)vy; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * qw = imatrix ? imatrix + QK_K*ibl : nullptr; + quantize_row_iq4_k_impl_bs16(QK_K, 16, x + QK_K*ibl, iq4 + ibl, + scales, weight, L, iq4k_values, qw, 7); + } + }; + auto row_size = ggml_row_size(GGML_TYPE_IQ4_K, n_per_row); + QHelper helper(imatrix, user_data, n_per_row, 16); + helper.quantize(nrows, src, dst, row_size, q_func); + return nrows * row_size; +} + +// +// ============================================== iq5_K +// +void dequantize_row_iq5_k(const block_iq5_k * x, float * y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint8_t * sl = x[i].scales_l; + const uint8_t * sh = x[i].scales_h; + + uint16_t extra = x[i].extra; + + int shift = 0; + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + + float dl1 = d * (((sl[2*ib64+0] & 0xf) | ((sh[ib64] << 4) & 0x30)) - 32); + float dl2 = d * (((sl[2*ib64+0] >> 4) | ((sh[ib64] << 2) & 0x30)) - 32); + float dl3 = d * (((sl[2*ib64+1] & 0xf) | ((sh[ib64] >> 0) & 0x30)) - 32); + float dl4 = d * (((sl[2*ib64+1] >> 4) | ((sh[ib64] >> 2) & 0x30)) - 32); + const int8_t * values1 = iq5nl_values + ((extra & 1) << 5); + const int8_t * values2 = iq5nl_values + ((extra & 2) << 4); + const int8_t * values3 = iq5nl_values + ((extra & 4) << 3); + const int8_t * values4 = iq5nl_values + ((extra & 8) << 2); + for (int j = 0; j < 16; ++j) { + y[j+ 0] = dl1 * values1[(qs[j+ 0] & 0xf) | (((qh[j+ 0] >> shift) & 1) << 4)]; + y[j+16] = dl2 * values2[(qs[j+16] & 0xf) | (((qh[j+16] >> shift) & 1) << 4)]; + y[j+32] = dl3 * values3[(qs[j+ 0] >> 4) | (((qh[j+ 0] >> shift) & 2) << 3)]; + y[j+48] = dl4 * values4[(qs[j+16] >> 4) | (((qh[j+16] >> shift) & 2) << 3)]; + } + y += 64; + qs += 32; + extra >>= 4; + shift += 2; + if (shift == 8) { qh += 32; shift = 0; } + } + + } +} + +void vec_dot_iq5_k_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ5_K, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + + const int nb = n / QK_K; + + const block_iq5_k * x = (const block_iq5_k *)vx; + const block_q8_K * y = (const block_q8_K *)vy; + + float sumf = 0; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d) * y[i].d; + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const uint8_t * sl = x[i].scales_l; + const uint8_t * sh = x[i].scales_h; + const int8_t * q8 = y[i].qs; + + uint16_t extra = x[i].extra; + + int shift = 0; + int sumb = 0; + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + + int dl1 = (((sl[2*ib64+0] & 0xf) | ((sh[ib64] << 4) & 0x30)) - 32); + int dl2 = (((sl[2*ib64+0] >> 4) | ((sh[ib64] << 2) & 0x30)) - 32); + int dl3 = (((sl[2*ib64+1] & 0xf) | ((sh[ib64] >> 0) & 0x30)) - 32); + int dl4 = (((sl[2*ib64+1] >> 4) | ((sh[ib64] >> 2) & 0x30)) - 32); + const int8_t * values1 = iq5nl_values + ((extra & 1) << 5); + const int8_t * values2 = iq5nl_values + ((extra & 2) << 4); + const int8_t * values3 = iq5nl_values + ((extra & 4) << 3); + const int8_t * values4 = iq5nl_values + ((extra & 8) << 2); + int sumi1 = 0, sumi2 = 0, sumi3 = 0, sumi4 = 0; + for (int j = 0; j < 16; ++j) { + sumi1 += q8[j+ 0] * values1[(qs[j+ 0] & 0xf) | (((qh[j+ 0] >> shift) & 1) << 4)]; + sumi2 += q8[j+16] * values2[(qs[j+16] & 0xf) | (((qh[j+16] >> shift) & 1) << 4)]; + sumi3 += q8[j+32] * values3[(qs[j+ 0] >> 4) | (((qh[j+ 0] >> shift) & 2) << 3)]; + sumi4 += q8[j+48] * values4[(qs[j+16] >> 4) | (((qh[j+16] >> shift) & 2) << 3)]; + } + sumb += dl1 * sumi1 + dl2 * sumi2 + dl3 * sumi3 + dl4 * sumi4; + q8 += 64; + qs += 32; + extra >>= 4; + shift += 2; + } + sumf += d * sumb; + + } + + *s = sumf; + +} + +namespace { +const int8_t iq5nl_index[248] = { + 0, 0, 0, 0, 0, 0, 32, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 33, 33, 2, 2, 2, 2, 2, 2, 2, 2, 2, 34, 34, 3, 3, + 3, 3, 3, 3, 3, 3, 35, 35, 4, 4, 4, 4, 4, 4, 4, 36, 36, 5, 5, 5, 5, 5, 5, 5, 37, 37, 6, 6, 6, 6, 6, 6, + 6, 38, 7, 7, 7, 7, 7, 7, 39, 39, 8, 8, 8, 8, 8, 40, 40, 9, 9, 9, 9, 9, 41, 41, 10, 10, 10, 10, 10, 42, 11, 11, + 11, 11, 11, 43, 12, 12, 12, 12, 12, 44, 13, 13, 13, 13, 13, 45, 14, 14, 14, 14, 14, 46, 15, 15, 15, 15, 47, 47, 16, 16, 16, 16, + 48, 17, 17, 17, 17, 17, 49, 18, 18, 18, 18, 18, 50, 19, 19, 19, 19, 19, 51, 20, 20, 20, 20, 20, 52, 21, 21, 21, 21, 21, 53, 53, + 22, 22, 22, 22, 22, 54, 54, 23, 23, 23, 23, 23, 23, 55, 24, 24, 24, 24, 24, 24, 24, 56, 25, 25, 25, 25, 25, 25, 25, 57, 57, 26, + 26, 26, 26, 26, 26, 26, 58, 58, 27, 27, 27, 27, 27, 27, 27, 27, 59, 28, 28, 28, 28, 28, 28, 28, 28, 28, 60, 29, 29, 29, 29, 29, + 29, 29, 29, 29, 29, 61, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 62, 31, 31, 31, 31, 31, 31 +}; +inline int best_index_iq5nl(const int8_t * values, float x) { + int ix = (int)x - values[0]; + if (ix < 0 || ix >= 247) return ix < 0 ? 0 : 31; + ix = iq5nl_index[ix]; + return ix < 32 ? ix : x - values[ix-32] < values[ix-31] - x ? ix-32 : ix-31; +} + +void quantize_row_iq5_k_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, + [[maybe_unused]] const quantize_user_data * user_data) { + const int ntry = 5; + const float step = 1.f; + + block_iq5_k * y = (block_iq5_k *)vy; + + float scales[QK_K/16]; + float weight[16]; + + const int8_t * shifted_values = iq5nl_values + 32; + + for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq5_k)); + y[ibl].d = GGML_FP32_TO_FP16(0.f); + + const float * xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 2*sumx2/QK_K; + + float max_scale = 0, max_abs_scale = 0; + uint16_t extra = 0; + + for (int ib = 0; ib < QK_K/16; ++ib) { + const float * xb = xbl + 16*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*16; + for (int j = 0; j < 16; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < 16; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < 16; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/iq5nl_values[0] : max/iq5nl_values[0]; + float id = 1/d; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq5nl(iq5nl_values, al); + float q = iq5nl_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq5nl(iq5nl_values, -al); + q = iq5nl_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + d = sumqx_p/sumq2_p; + float best = d*sumqx_p; + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d*sumqx_m; + } + bool is_shifted = false; + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry*step + iq5nl_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq5nl(iq5nl_values, al); + float q = iq5nl_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq5nl(iq5nl_values, -al); + q = iq5nl_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = false; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = false; + } + id = (itry*step + shifted_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq5nl(shifted_values, al); + float q = shifted_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq5nl(shifted_values, -al); + q = shifted_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = true; + } + } + if (d) { + const int8_t * block_values = is_shifted ? shifted_values : iq5nl_values; + float sumqx = 0, sumq2 = 0; + id = 1/d; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq5nl(block_values, al); + float q = block_values[l]; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + if (sumq2 > 0) d = sumqx/sumq2; + } + scales[ib] = d; + if (is_shifted) extra |= (1 << ib); + + float abs_scale = fabsf(scales[ib]); + if (abs_scale > max_abs_scale) { + max_abs_scale = abs_scale; max_scale = scales[ib]; + } + + } + + if (!max_abs_scale) continue; + float d = -max_scale/32; + y[ibl].d = GGML_FP32_TO_FP16(d); + y[ibl].extra = extra; + + float id = 1/d; + + float sumqx = 0, sumq2 = 0; + for (int ib = 0; ib < QK_K/16; ++ib) { + int ls = nearest_int(id*scales[ib]); + ls = MAX(-32, MIN(31, ls)); + int uls = ls + 32; + y[ibl].scales_l[ib/2] |= ((uls & 0xf) << 4*(ib%2)); + y[ibl].scales_h[ib/4] |= ((uls >> 4) << 2*(ib%4)); + float dl = d * ls; + if (dl) { + const int8_t * block_values = y[ibl].extra & (1 << ib) ? shifted_values : iq5nl_values; + const float * xb = xbl + 16*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*16; + for (int j = 0; j < 16; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < 16; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float idl = 1/dl; + int ib32 = ib/2; + int offset = 16*(ib%2); + uint8_t * qs = y[ibl].qs + 32*(ib32/2) + offset; + uint8_t * qh = y[ibl].qh + 32*(ib32/8) + offset; + for (int j = 0; j < 16; ++j) { + const float al = idl*xb[j]; + int ibest = best_index_iq5nl(block_values, al); + qs[j] |= ((ibest & 0xf) << 4*(ib32%2)); + qh[j] |= ((ibest >> 4) << (ib32%8)); + float w = weight[j]; + float q = block_values[ibest]*ls; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + } + } + if (sumq2 > 0) y[ibl].d = GGML_FP32_TO_FP16(sumqx/sumq2); + + } + +} + +} + +void quantize_row_iq5_k_ref(const float * x, block_iq5_k * y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq5_k(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq5_k(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + block_iq5_k * y = (block_iq5_k *)vy; + quantize_row_iq5_k_ref(x, y, k); +} + +size_t quantize_iq5_k(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + QHelper helper(imatrix, user_data, n_per_row, 16); + auto row_size = ggml_row_size(GGML_TYPE_IQ5_K, n_per_row); + helper.quantize(nrows, src, dst, row_size, quantize_row_iq5_k_impl); + return nrows * row_size; +} + +// +// ============================================== iq6_K +// +#define A_IQ6K -127.f +#define B_IQ6K 6.2568f +#define C_IQ6K 0.11218f +#define D_IQ6K 0.0011972f +#define S_IQ6K 1.f + +void dequantize_row_iq6_k(const block_iq6_k * x, float * y, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + + for (int i = 0; i < nb; i++) { + + const float d = GGML_FP16_TO_FP32(x[i].d); + const uint8_t * qs = x[i].qs; + const uint8_t * qh = x[i].qh; + const int8_t * sl = x[i].scales; + + uint16_t extra = x[i].extra; + + int shift = 0; + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + + float dl1 = d * sl[4*ib64 + 0]; + float dl2 = d * sl[4*ib64 + 1]; + float dl3 = d * sl[4*ib64 + 2]; + float dl4 = d * sl[4*ib64 + 3]; + float m1 = extra & 1 ? S_IQ6K : 0; + float m2 = extra & 2 ? S_IQ6K : 0; + float m3 = extra & 4 ? S_IQ6K : 0; + float m4 = extra & 8 ? S_IQ6K : 0; + for (int j = 0; j < 16; ++j) { + float q1 = ((qs[j+ 0] & 0xf) | (((qh[j+ 0] >> shift) & 0x03) << 4)); + float q2 = ((qs[j+16] & 0xf) | (((qh[j+16] >> shift) & 0x03) << 4)); + float q3 = ((qs[j+ 0] >> 4) | (((qh[j+ 0] >> shift) & 0x0c) << 2)); + float q4 = ((qs[j+16] >> 4) | (((qh[j+16] >> shift) & 0x0c) << 2)); + y[j+ 0] = dl1 * (A_IQ6K + q1*(B_IQ6K + q1*(-C_IQ6K + q1*D_IQ6K)) + m1); + y[j+16] = dl2 * (A_IQ6K + q2*(B_IQ6K + q2*(-C_IQ6K + q2*D_IQ6K)) + m2); + y[j+32] = dl3 * (A_IQ6K + q3*(B_IQ6K + q3*(-C_IQ6K + q3*D_IQ6K)) + m3); + y[j+48] = dl4 * (A_IQ6K + q4*(B_IQ6K + q4*(-C_IQ6K + q4*D_IQ6K)) + m4); + } + y += 64; + qs += 32; + extra >>= 4; + shift += 4; + if (shift == 8) { qh += 32; shift = 0; } + } + + } +} + +void vec_dot_iq6_k_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ6_K, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + + GGML_ABORT("not implemented"); + + // TODO + //const int nb = n / QK_K; + + //const block_iq5_k * x = (const block_iq5_k *)vx; + //const block_q8_K * y = (const block_q8_K *)vy; + + //float sumf = 0; + + //for (int i = 0; i < nb; i++) { + + // const float d = GGML_FP16_TO_FP32(x[i].d) * y[i].d; + // const uint8_t * qs = x[i].qs; + // const uint8_t * qh = x[i].qh; + // const uint8_t * sl = x[i].scales_l; + // const uint8_t * sh = x[i].scales_h; + // const int8_t * q8 = y[i].qs; + + // uint16_t extra = x[i].extra; + + // int shift = 0; + // int sumb = 0; + // for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + + // int dl1 = (((sl[2*ib64+0] & 0xf) | ((sh[ib64] << 4) & 0x30)) - 32); + // int dl2 = (((sl[2*ib64+0] >> 4) | ((sh[ib64] << 2) & 0x30)) - 32); + // int dl3 = (((sl[2*ib64+1] & 0xf) | ((sh[ib64] >> 0) & 0x30)) - 32); + // int dl4 = (((sl[2*ib64+1] >> 4) | ((sh[ib64] >> 2) & 0x30)) - 32); + // const int8_t * values1 = iq5nl_values + ((extra & 1) << 5); + // const int8_t * values2 = iq5nl_values + ((extra & 2) << 4); + // const int8_t * values3 = iq5nl_values + ((extra & 4) << 3); + // const int8_t * values4 = iq5nl_values + ((extra & 8) << 2); + // int sumi1 = 0, sumi2 = 0, sumi3 = 0, sumi4 = 0; + // for (int j = 0; j < 16; ++j) { + // sumi1 += q8[j+ 0] * values1[(qs[j+ 0] & 0xf) | (((qh[j+ 0] >> shift) & 1) << 4)]; + // sumi2 += q8[j+16] * values2[(qs[j+16] & 0xf) | (((qh[j+16] >> shift) & 1) << 4)]; + // sumi3 += q8[j+32] * values3[(qs[j+ 0] >> 4) | (((qh[j+ 0] >> shift) & 2) << 3)]; + // sumi4 += q8[j+48] * values4[(qs[j+16] >> 4) | (((qh[j+16] >> shift) & 2) << 3)]; + // } + // sumb += dl1 * sumi1 + dl2 * sumi2 + dl3 * sumi3 + dl4 * sumi4; + // q8 += 64; + // qs += 32; + // extra >>= 4; + // shift += 2; + // } + // sumf += d * sumb; + + //} + + //*s = sumf; + +} + +namespace { + +inline int best_index(int n, const float * val, float x) { + if (x <= val[0]) return 0; + if (x >= val[n-1]) return n-1; + int ml = 0, mu = n-1; + while (mu-ml > 1) { + int mav = (ml+mu)/2; + if (x < val[mav]) mu = mav; else ml = mav; + } + return x - val[mu-1] < val[mu] - x ? mu-1 : mu; +} +uint8_t iq6nl_index[249] = { + 0, 0, 0, 64, 1, 1, 1, 1, 1, 65, 2, 2, 2, 2, 2, 66, 3, 3, 3, 3, 67, 67, 4, 4, 4, 4, 68, 5, 5, 5, 5, 69, + 69, 6, 6, 6, 70, 70, 7, 7, 7, 71, 8, 8, 8, 72, 72, 9, 9, 9, 73, 73, 10, 10, 10, 74, 11, 11, 11, 75, 12, 12, 12, 76, + 13, 13, 13, 77, 14, 14, 14, 78, 15, 15, 79, 79, 16, 16, 80, 17, 17, 81, 81, 18, 18, 82, 19, 19, 83, 83, 20, 84, 84, 21, 85, 85, + 22, 86, 86, 23, 87, 87, 24, 88, 88, 25, 89, 89, 26, 90, 90, 27, 91, 91, 28, 92, 29, 93, 93, 30, 94, 94, 31, 95, 95, 32, 96, 33, + 97, 97, 34, 98, 98, 35, 99, 99, 36, 100, 100, 37, 101, 38, 102, 102, 39, 103, 103, 40, 104, 104, 41, 41, 105, 42, 42, 106, 106, 43, 107, 107, + 44, 108, 108, 45, 45, 109, 46, 46, 46, 110, 47, 47, 111, 111, 48, 48, 112, 49, 49, 49, 113, 50, 50, 50, 114, 51, 51, 51, 115, 52, 52, 52, + 116, 116, 53, 53, 53, 117, 54, 54, 54, 118, 118, 55, 55, 55, 119, 119, 56, 56, 56, 120, 120, 57, 57, 57, 121, 121, 58, 58, 58, 58, 122, 59, + 59, 59, 59, 123, 123, 60, 60, 60, 60, 124, 61, 61, 61, 61, 61, 125, 62, 62, 62, 62, 62, 126, 63, 63, 63, +}; +inline int best_index_iq6nl(const float * values, float x) { + int ix = (int)(x - values[0]); + if (ix < 0 || ix >= 249) return ix < 0 ? 0 : 63; + ix = iq6nl_index[ix]; + return ix < 64 ? ix : x - values[ix-64] < values[ix-63] - x ? ix-64 : ix-63; + //if (x <= val[0]) return 0; + //if (x >= val[63]) return 63; + //int index = iq6nl_index[int(x - val[0])]; + //return index < 64 ? index : x - val[index-64] < val[index-63] - x ? index - 64 : index - 63; +} + + +void quantize_row_iq6_k_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, const float * values, const float * shifted_values) { + const int ntry = 5; + const float step = 1.f; + + block_iq6_k * y = (block_iq6_k *)vy; + + float scales[QK_K/16]; + float weight[16]; + + for (int ibl = 0; ibl < n_per_row/QK_K; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq6_k)); + y[ibl].d = GGML_FP32_TO_FP16(0.f); + + const float * xbl = x + ibl*QK_K; + float sumx2 = 0; + for (int j = 0; j < QK_K; ++j) sumx2 += xbl[j]*xbl[j]; + const float sigma2 = 2*sumx2/QK_K; + + float max_scale = 0, max_abs_scale = 0; + uint16_t extra = 0; + + for (int ib = 0; ib < QK_K/16; ++ib) { + const float * xb = xbl + 16*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*16; + for (int j = 0; j < 16; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < 16; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < 16; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/values[0] : max/values[0]; + float id = 1/d; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + //int l = best_index(64, values, al); + int l = best_index_iq6nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + //l = best_index(64, values, -al); + l = best_index_iq6nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + d = sumqx_p/sumq2_p; + float best = d*sumqx_p; + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d*sumqx_m; + } + bool is_shifted = false; + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry*step + values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + //int l = best_index(64, values, al); + int l = best_index_iq6nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + //l = best_index(64, values, -al); + l = best_index_iq6nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = false; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = false; + } + id = (itry*step + shifted_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + //int l = best_index(64, shifted_values, al); + int l = best_index_iq6nl(shifted_values, al); + float q = shifted_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + //l = best_index(64, shifted_values, -al); + l = best_index_iq6nl(shifted_values, -al); + q = shifted_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = true; + } + } + if (d) { + const float * block_values = is_shifted ? shifted_values : values; + float sumqx = 0, sumq2 = 0; + id = 1/d; + for (int j = 0; j < 16; ++j) { + float w = weight[j]; + float al = id*xb[j]; + //int l = best_index(64, block_values, al); + int l = best_index_iq6nl(block_values, al); + float q = block_values[l]; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + if (sumq2 > 0) d = sumqx/sumq2; + } + scales[ib] = d; + if (is_shifted) extra |= (1 << ib); + + float abs_scale = fabsf(scales[ib]); + if (abs_scale > max_abs_scale) { + max_abs_scale = abs_scale; max_scale = scales[ib]; + } + + } + + if (!max_abs_scale) continue; + float d = -max_scale/127; + y[ibl].d = GGML_FP32_TO_FP16(d); + y[ibl].extra = extra; + + float id = 1/d; + + float sumqx = 0, sumq2 = 0; + for (int ib = 0; ib < QK_K/16; ++ib) { + int ls = nearest_int(id*scales[ib]); + ls = MAX(-127, MIN(127, ls)); + y[ibl].scales[ib] |= ls; + float dl = d * ls; + if (dl) { + const float * block_values = y[ibl].extra & (1 << ib) ? shifted_values : values; + const float * xb = xbl + 16*ib; + if (quant_weights) { + const float * qw = quant_weights + ibl*QK_K + ib*16; + for (int j = 0; j < 16; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < 16; ++j) weight[j] = 0.25f*sigma2 + xb[j]*xb[j]; + } + float idl = 1/dl; + int ib32 = ib/2; + int offset = 16*(ib%2); + uint8_t * qs = y[ibl].qs + 32*(ib32/2) + offset; + uint8_t * qh = y[ibl].qh + 32*(ib32/4) + offset; + for (int j = 0; j < 16; ++j) { + const float al = idl*xb[j]; + //int ibest = best_index(64, block_values, al); + int ibest = best_index_iq6nl(block_values, al); + qs[j] |= ((ibest & 0xf) << 4*(ib32%2)); + qh[j] |= ((ibest >> 4) << 2*(ib32%4)); + float w = weight[j]; + float q = block_values[ibest]*ls; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + } + } + if (sumq2 > 0) y[ibl].d = GGML_FP32_TO_FP16(sumqx/sumq2); + + } +} + +} + +void quantize_row_iq6_k_ref(const float * x, block_iq6_k * y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq6_k(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq6_k(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + block_iq6_k * y = (block_iq6_k *)vy; + quantize_row_iq6_k_ref(x, y, k); +} + +size_t quantize_iq6_k(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + float values[128]; + for (int i = 0; i < 64; ++i) { + values[i] = iq6nl_values[i]; + values[i+64] = values[i] + S_IQ6K; + } + auto q_func = [values] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + quantize_row_iq6_k_impl(x, vy, n_per_row, imatrix, values, values + 64); + }; + auto row_size = ggml_row_size(GGML_TYPE_IQ6_K, n_per_row); + QHelper helper(imatrix, user_data, n_per_row, 16); + helper.quantize(nrows, src, dst, row_size, q_func); + return nrows * row_size; +} + +namespace { +template +void iqk_quantize_row_q8_K_T(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + const int nb = k / QK_K; + iqk_block_q8_K * y = (iqk_block_q8_K *)vy; // opencoti F5-opt W2 (#290): ik layout has extra `sum` +#ifdef __AVX2__ + const __m256 signBit = _mm256_set1_ps(-0.0f); + const __m256i perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7); + for (int i = 0; i < nb; i++) { + const float * xb = x + i*QK_K; + __m256 maxAbs = _mm256_setzero_ps(); + const float * xx = xb; + for (int ib = 0; ib < QK_K/8; ++ib) { + const __m256 v = _mm256_loadu_ps(xx); xx += 8; + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps(signBit, v)); + } + const float maxScalar = hmax_f32_8(maxAbs); + const float d = maxScalar / 127.f; + y[i].d = d; + const float id = ( maxScalar != 0.0f ) ? 127.f / maxScalar : 0.0f; + const __m256 mul = _mm256_set1_ps( id ); + xx = xb; + int8_t * q8 = y[i].qs; + int block_sum_i32 = 0; + float block_sum_f32 = 0; + for (int ib = 0; ib < QK_K/32; ++ib) { + __m256 v0 = _mm256_mul_ps(mul, _mm256_loadu_ps(xx)); xx += 8; + __m256 v1 = _mm256_mul_ps(mul, _mm256_loadu_ps(xx)); xx += 8; + __m256 v2 = _mm256_mul_ps(mul, _mm256_loadu_ps(xx)); xx += 8; + __m256 v3 = _mm256_mul_ps(mul, _mm256_loadu_ps(xx)); xx += 8; + v0 = _mm256_round_ps(v0, _MM_ROUND_NEAREST); + v1 = _mm256_round_ps(v1, _MM_ROUND_NEAREST); + v2 = _mm256_round_ps(v2, _MM_ROUND_NEAREST); + v3 = _mm256_round_ps(v3, _MM_ROUND_NEAREST); + __m256i i0 = _mm256_cvtps_epi32(v0); + __m256i i1 = _mm256_cvtps_epi32(v1); + __m256i i2 = _mm256_cvtps_epi32(v2); + __m256i i3 = _mm256_cvtps_epi32(v3); + if constexpr (q8_type == 1) { + int bsum = hsum_i32_8(_mm256_add_epi32(_mm256_add_epi32(i0, i1), _mm256_add_epi32(i2, i3))); + auto bs = (float *)y[i].bsums; + bs[ib] = d*bsum; + block_sum_f32 += bs[ib]; + } else { + y[i].bsums[2*ib+0] = hsum_i32_8(_mm256_add_epi32(i0, i1)); + y[i].bsums[2*ib+1] = hsum_i32_8(_mm256_add_epi32(i2, i3)); + block_sum_i32 += y[i].bsums[2*ib+0] + y[i].bsums[2*ib+1]; + } + i0 = _mm256_packs_epi32( i0, i1 ); + i2 = _mm256_packs_epi32( i2, i3 ); + i0 = _mm256_packs_epi16( i0, i2 ); + i0 = _mm256_permutevar8x32_epi32( i0, perm ); + _mm256_storeu_si256((__m256i *)q8, i0); + q8 += 32; + } + if constexpr (q8_type == 1) { + y[i].sum = block_sum_f32; + } else { + y[i].sum = d*block_sum_i32; + } + //if constexpr (q8_type == 2) { + // auto bs = (float *)y[i].bsums; + // float sum = 0; + // for (int ib = 0; ib < QK_K/32; ++ib) sum += bs[ib]; + // bs[0] = sum; + //} + } +#else + for (int i = 0; i < nb; i++) { + + float max = 0; + float amax = 0; + for (int j = 0; j < QK_K; ++j) { + float ax = fabsf(x[j]); + if (ax > amax) { + amax = ax; max = x[j]; + } + } + if (!amax) { + y[i].d = 0; + memset(y[i].qs, 0, QK_K); + x += QK_K; + continue; + } + //const float iscale = -128.f/max; + // We need this change for IQ2_XXS, else the AVX implementation becomes very awkward + const float iscale = -127.f/max; + for (int j = 0; j < QK_K; ++j) { + int v = nearest_int(iscale*x[j]); + y[i].qs[j] = MIN(127, v); + } + float d = 1/iscale; + if constexpr (q8_type == 1) { + auto bs = (float *)y[i].bsums; + float sum = 0; + for (int j = 0; j < QK_K/32; ++j) { + int sum = 0; + for (int ii = 0; ii < 32; ++ii) { + sum += y[i].qs[j*32 + ii]; + } + bs[j] = d*sum; + sum += bs[j]; + } + y[i].sum = sum; + } else { + int tot = 0; + for (int j = 0; j < QK_K/16; ++j) { + int sum = 0; + for (int ii = 0; ii < 16; ++ii) { + sum += y[i].qs[j*16 + ii]; + } + y[i].bsums[j] = sum; + tot += sum; + } + y[i].sum = d*tot; + } + y[i].d = d; + x += QK_K; + } +#endif +} +} + +void iqk_quantize_row_q8_K(const float * x, void * vy, int64_t k) { + iqk_quantize_row_q8_K_T<0>(x, vy, k); +} + +void quantize_row_q8_K32(const float * x, void * vy, int64_t k) { + iqk_quantize_row_q8_K_T<1>(x, vy, k); +} + +void quantize_row_q8_KR8(const float * x, void * vy, int64_t k) { + iqk_quantize_row_q8_K_T<2>(x, vy, k); +} + +namespace { +// TODO: merge this with the above template +void iqk_quantize_row_q8_K128(const float * x, void * vy, int64_t k) { + constexpr int kBlockSize = 128; + assert(k % kBlockSize == 0); + const int nb = k / kBlockSize; + auto y = (block_q8_K128 *)vy; +#ifdef __AVX2__ + const __m256 signBit = _mm256_set1_ps(-0.0f); + const __m256i perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7); + for (int i = 0; i < nb; i++) { + const float * xb = x + i*kBlockSize; + __m256 maxAbs = _mm256_setzero_ps(); + const float * xx = xb; + for (int ib = 0; ib < kBlockSize/8; ++ib) { + const __m256 v = _mm256_loadu_ps(xx); xx += 8; + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps(signBit, v)); + } + const float maxScalar = hmax_f32_8(maxAbs); + const float d = maxScalar / 127.f; + y[i].d = d; + const float id = ( maxScalar != 0.0f ) ? 127.f / maxScalar : 0.0f; + const __m256 mul = _mm256_set1_ps( id ); + xx = xb; + int8_t * q8 = y[i].qs; + for (int ib = 0; ib < kBlockSize/32; ++ib) { + __m256 v0 = _mm256_mul_ps(mul, _mm256_loadu_ps(xx)); xx += 8; + __m256 v1 = _mm256_mul_ps(mul, _mm256_loadu_ps(xx)); xx += 8; + __m256 v2 = _mm256_mul_ps(mul, _mm256_loadu_ps(xx)); xx += 8; + __m256 v3 = _mm256_mul_ps(mul, _mm256_loadu_ps(xx)); xx += 8; + v0 = _mm256_round_ps(v0, _MM_ROUND_NEAREST); + v1 = _mm256_round_ps(v1, _MM_ROUND_NEAREST); + v2 = _mm256_round_ps(v2, _MM_ROUND_NEAREST); + v3 = _mm256_round_ps(v3, _MM_ROUND_NEAREST); + __m256i i0 = _mm256_cvtps_epi32(v0); + __m256i i1 = _mm256_cvtps_epi32(v1); + __m256i i2 = _mm256_cvtps_epi32(v2); + __m256i i3 = _mm256_cvtps_epi32(v3); + y[i].bsums[ib] = hsum_i32_8(_mm256_add_epi32(_mm256_add_epi32(i0, i1), _mm256_add_epi32(i2, i3))); + i0 = _mm256_packs_epi32( i0, i1 ); + i2 = _mm256_packs_epi32( i2, i3 ); + i0 = _mm256_packs_epi16( i0, i2 ); + i0 = _mm256_permutevar8x32_epi32( i0, perm ); + _mm256_storeu_si256((__m256i *)q8, i0); + q8 += 32; + } + } +#elif defined __ARM_NEON + int32x4_t ival[8]; + for (int i = 0; i < nb; i++) { + const float * xb = x + i*kBlockSize; + auto vmax = vdupq_n_f32(0.f); + for (int j = 0; j < kBlockSize; j += 4) { + vmax = vmaxq_f32(vmax, vabsq_f32(vld1q_f32(xb + j))); + } + auto smax = vmaxvq_f32(vmax); + if (!smax) { + std::memset(&y[i], 0, sizeof(y[i])); + continue; + } + y[i].d = smax/127; + auto vid = vdupq_n_f32(127/smax); + for (int ib = 0; ib < kBlockSize/32; ++ib) { + auto isum = vdupq_n_s32(0); + for (int k = 0; k < 8; ++k) { + auto val = vld1q_f32(xb + 32*ib + 4*k); + ival[k] = vcvtnq_s32_f32(vmulq_f32(val, vid)); + isum = vaddq_s32(isum, ival[k]); + } + y[i].bsums[ib] = vaddvq_s32(isum); + for (int k = 0; k < 4; ++k) { + auto i16 = vcombine_s16(vmovn_s32(ival[2*k+0]), vmovn_s32(ival[2*k+1])); + vst1_s8(y[i].qs + 32*ib + 8*k, vmovn_s16(i16)); + } + } + } +#else + for (int i = 0; i < nb; i++) { + + float amax = 0; + for (int j = 0; j < kBlockSize; ++j) { + float ax = std::abs(x[j]); + amax = std::max(amax, ax); + } + if (!amax) { + y[i].d = 0; + memset(y[i].qs, 0, kBlockSize); + memset(y[i].bsums, 0, kBlockSize/32*(sizeof(int16_t))); + x += kBlockSize; + continue; + } + const float iscale = 127.f/amax; + for (int j = 0; j < kBlockSize; ++j) { + int v = nearest_int(iscale*x[j]); + y[i].qs[j] = v; + } + for (int j = 0; j < kBlockSize/32; ++j) { + int sum = 0; + for (int ii = 0; ii < 32; ++ii) { + sum += y[i].qs[j*32 + ii]; + } + y[i].bsums[j] = sum; + } + y[i].d = 1/iscale; + x += kBlockSize; + } +#endif +} +// TODO: merge this with the above template +void iqk_quantize_row_q8_KV(const float * x, void * vy, int64_t k) { + assert(k % 32 == 0); + auto dptr = (float *)vy; + auto q8 = (int8_t *)(dptr + 2); +#ifdef __AVX2__ + const __m256 signBit = _mm256_set1_ps(-0.0f); + const __m256i perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7); + __m256 maxAbs = _mm256_setzero_ps(); + for (int ib = 0; ib < k/8; ++ib) { + const __m256 v = _mm256_loadu_ps(x + 8*ib); + maxAbs = _mm256_max_ps( maxAbs, _mm256_andnot_ps(signBit, v)); + } + const float maxScalar = hmax_f32_8(maxAbs); + if (!maxScalar) { + dptr[0] = dptr[1] = 0; + std::memset(q8, 0, k*sizeof(int8_t)); + return; + } + dptr[0] = maxScalar / 127.f; + auto mul = _mm256_set1_ps(1/dptr[0]); + auto isum = _mm256_setzero_si256(); + for (int i = 0; i < k/32; i++) { + __m256 v0 = _mm256_mul_ps(mul, _mm256_loadu_ps(x + 32*i + 0)); + __m256 v1 = _mm256_mul_ps(mul, _mm256_loadu_ps(x + 32*i + 8)); + __m256 v2 = _mm256_mul_ps(mul, _mm256_loadu_ps(x + 32*i + 16)); + __m256 v3 = _mm256_mul_ps(mul, _mm256_loadu_ps(x + 32*i + 24)); + v0 = _mm256_round_ps(v0, _MM_ROUND_NEAREST); + v1 = _mm256_round_ps(v1, _MM_ROUND_NEAREST); + v2 = _mm256_round_ps(v2, _MM_ROUND_NEAREST); + v3 = _mm256_round_ps(v3, _MM_ROUND_NEAREST); + __m256i i0 = _mm256_cvtps_epi32(v0); + __m256i i1 = _mm256_cvtps_epi32(v1); + __m256i i2 = _mm256_cvtps_epi32(v2); + __m256i i3 = _mm256_cvtps_epi32(v3); + isum = _mm256_add_epi32(isum, _mm256_add_epi32(_mm256_add_epi32(i0, i1), _mm256_add_epi32(i2, i3))); + i0 = _mm256_packs_epi32( i0, i1 ); + i2 = _mm256_packs_epi32( i2, i3 ); + i0 = _mm256_packs_epi16( i0, i2 ); + i0 = _mm256_permutevar8x32_epi32( i0, perm ); + _mm256_storeu_si256((__m256i *)q8, i0); + q8 += 32; + } + auto iptr = (int32_t *)(dptr + 1); + iptr[0] = hsum_i32_8(isum); +#elif defined __ARM_NEON + int32x4_t ival[8]; + auto vmax = vdupq_n_f32(0.f); + for (int j = 0; j < k; j += 4) { + vmax = vmaxq_f32(vmax, vabsq_f32(vld1q_f32(x + j))); + } + auto smax = vmaxvq_f32(vmax); + if (!smax) { + dptr[0] = dptr[1] = 0; + std::memset(q8, 0, k*sizeof(int8_t)); + return; + } + dptr[0] = smax/127; + auto vid = vdupq_n_f32(1/dptr[0]); + auto isum = vdupq_n_s32(0); + for (int ib = 0; ib < k/32; ++ib) { + auto xb = x + 32*ib; + for (int k = 0; k < 8; ++k) { + auto val = vld1q_f32(xb + 4*k); + ival[k] = vcvtnq_s32_f32(vmulq_f32(val, vid)); + isum = vaddq_s32(isum, ival[k]); + } + for (int k = 0; k < 4; ++k) { + auto i16 = vcombine_s16(vmovn_s32(ival[2*k+0]), vmovn_s32(ival[2*k+1])); + vst1_s8(q8, vmovn_s16(i16)); + q8 += 8; + } + } + auto iptr = (int32_t *)(dptr + 1); + iptr[0] = vaddvq_s32(isum); +#else + float amax = 0; + for (int j = 0; j < k; ++j) { + float ax = std::abs(x[j]); + amax = std::max(amax, ax); + } + if (!amax) { + dptr[0] = dptr[1] = 0; + std::memset(q8, 0, k*sizeof(int8_t)); + return; + } + dptr[0] = amax/127; + float id = 1/dptr[0]; + int isum = 0; + for (int i = 0; i < k; i++) { + q8[i] = nearest_int(id*x[i]); + isum += q8[i]; + } + auto iptr = (int32_t *)(dptr + 1); + iptr[0] = isum; +#endif +} +} + +void quantize_row_q8_K128(const float * x, void * vy, int64_t k) { + iqk_quantize_row_q8_K128(x, vy, k); +} + +// ============================== MXFP4 + +namespace { +inline int best_index_mxfp4(float d, const int8_t * values, float x) { + float best = std::abs(x - d*values[0]); + int index = 0; + for (int j = 1; j < 16; ++j) { + float diff = std::abs(x - d*values[j]); + if (diff < best) { best = diff; index = j; } + } + return index; +} +static void quantize_row_mxfp4_impl(int n_per_row, const float * x, char * cy, + [[maybe_unused]] float * weight, + const int8_t * values, + [[maybe_unused]] const float * quant_weights, + [[maybe_unused]] const int ntry) { + + GGML_ASSERT(n_per_row % QK_MXFP4 == 0); + GGML_UNUSED(quant_weights); + + block_mxfp4 * y = (block_mxfp4 *)cy; + + //int last_ibl = -1; + //float sigma2 = 0; + + //const uint8_t e = (uint8_t) (floorf(log2f(amax)) - 2 + 127); + // -> log2f(amax) ~ e - 125 -> amax = 2^(e - 125) + //const float d = GGML_E8M0_TO_FP32_HALF(e); + + for (int ib = 0; ib < n_per_row/QK_MXFP4; ++ib) { + memset(&y[ib], 0, sizeof(block_mxfp4)); + const float * xb = x + ib*QK_MXFP4; + //if (int ibl = ib/(QK_K/QK_MXFP4); ibl != last_ibl) { + // int n = std::min(QK_K, n_per_row - ib*QK_MXFP4); + // float sumx2 = 0; + // for (int j = 0; j < n; ++j) sumx2 += xb[j]*xb[j]; + // sigma2 = 2.0f*sumx2/n; + // last_ibl = ibl; + //} + //if (quant_weights) { + // const float * qw = quant_weights + ib*QK_MXFP4; + // for (int j = 0; j < QK_MXFP4; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + //} else { + // for (int j = 0; j < QK_MXFP4; ++j) weight[j] = xb[j]*xb[j]; + //} + float amax = 0; + for (int j = 0; j < QK_MXFP4; ++j) { + float ax = fabsf(xb[j]); + amax = std::max(amax, ax); + } + if (!amax) { + continue; + } + const uint8_t e = (uint8_t) (floorf(log2f(amax)) - 2 + 127); + const float d = GGML_E8M0_TO_FP32_HALF(e); + y[ib].e = e; + for (int j = 0; j < QK_MXFP4/2; ++j) { + uint8_t v0 = best_index_mxfp4(d, values, xb[j]); + uint8_t v1 = best_index_mxfp4(d, values, xb[j+QK_MXFP4/2]); + y[ib].qs[j] = v0 | (v1 << 4); + } + } +} +} + +void iqk_quantize_row_mxfp4_ref(const float * x, block_mxfp4 * y, int64_t k) { + iqk_quantize_mxfp4(x, (void *)y, 1, k, nullptr, nullptr); +} + +void iqk_quantize_row_mxfp4(const float * x, void * y, int64_t k) { + iqk_quantize_mxfp4(x, (void *)y, 1, k, nullptr, nullptr); +} + +size_t iqk_quantize_mxfp4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + constexpr int kBlockSize = QK_MXFP4; + GGML_ASSERT(n_per_row%kBlockSize == 0); + auto row_size = ggml_row_size(GGML_TYPE_MXFP4, n_per_row); + char * qrow = (char *)dst; + float weight[kBlockSize]; + for (int64_t row = 0; row < nrows; ++row) { + quantize_row_mxfp4_impl(n_per_row, src, qrow, weight, kvalues_mxfp4, imatrix, 7); + src += n_per_row; + qrow += row_size; + } + return nrows * row_size; +} + +void iqk_dequantize_row_mxfp4(const block_mxfp4 * x, float * y, int64_t k) { + constexpr int kBlockSize = QK_MXFP4; + GGML_ASSERT(k%kBlockSize == 0); + int nblock = k/kBlockSize; + for (int ib = 0; ib < nblock; ++ib) { + float d = GGML_E8M0_TO_FP32_HALF(x[ib].e); + for (int j = 0; j < kBlockSize/2; ++j) { + y[j ] = d * kvalues_mxfp4[x[ib].qs[j] & 0xf]; + y[j+kBlockSize/2] = d * kvalues_mxfp4[x[ib].qs[j] >> 4]; + } + y += kBlockSize; + } +} + +void vec_dot_mxfp4_q8_0_x4(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_MXFP4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK_MXFP4 == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + //const block_mxfp4 * x = (const block_mxfp4 *)vx; + //const block_q8_K * y = (const block_q8_K *)vy; + //int nblock = n/QK_MXFP4; + //float sumf = 0; + //for (int ibl = 0; ibl < nblock; ++ibl) { + // //int sumi = 0; + // auto qy = y[ibl].qs; + // auto qx = x[ibl].qs; + // float db = d * y[ibl].d; + // for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + // float dl = db * ((x[ibl].scales[ib] & 254) - 127); + // //int ls = (x[ibl].scales[ib] & 254) - 127; + // const int8_t * values = iq4k_values + ((x[ibl].scales[ib] & 1) << 4); + // int suml = 0; + // for (int j = 0; j < kBlockSize/2; ++j) { + // suml += qy[j ] * values[qx[j] & 0xf] + // + qy[j + kBlockSize/2] * values[qx[j] >> 4]; + // } + // sumf += dl * suml; + // //sumi += ls * suml; + // qy += kBlockSize; + // qx += kBlockSize/2; + // } + // //sumf += d * y[ibl].d * sumi; + //} + //*s = sumf; +} + +namespace { +static void quantize_row_iq4_k_impl_bs128(const int super_block_size, const int block_size, + int n_per_row, const float * x, char * cy, + float * all_scales, float * weight, + const int8_t * values, + const float * quant_weights, + const int ntry) { + + //GGML_ASSERT(super_block_size == 256 && block_size == 128); + + float * dptr = (float *)cy; + block_iq4_ks * y = (block_iq4_ks *)(dptr + 1); + + const int8_t * shifted_values = values + 16; + + float amax_scale = 0; + + for (int ibl = 0; ibl < n_per_row/super_block_size; ++ibl) { + memset(&y[ibl], 0, sizeof(block_iq4_ks)); + const float * xbl = x + ibl*super_block_size; + auto scales = all_scales + ibl*(super_block_size/block_size); + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= 2.f/super_block_size; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const float * xb = xbl + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ibl*super_block_size + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < block_size; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/values[0] : max/values[0]; + float id = 1/d; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq4nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq4nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + d = sumqx_p/sumq2_p; + bool is_shifted = false; + float best = d*sumqx_p; + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d*sumqx_m; + } + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry + values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq4nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq4nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = false; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = false; + } + id = (itry + shifted_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq4nl(shifted_values, al); + float q = shifted_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq4nl(shifted_values, -al); + q = shifted_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = true; + } + } + if (is_shifted) y[ibl].scales[ib] = 0x01; + scales[ib] = d; + amax_scale = std::max(amax_scale, std::abs(d)); + } + } + float d = amax_scale/127; + *dptr = d; + if (!d) return; + float id = d ? 1/d : 0.f; + float sumqx = 0, sumq2 = 0; + //float mse = 0; + for (int ibl = 0; ibl < n_per_row/super_block_size; ++ibl) { + const float * xbl = x + ibl*super_block_size; + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= 2.f/super_block_size; + auto scales = all_scales + (super_block_size/block_size)*ibl; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const int8_t * block_values = y[ibl].scales[ib] & 0x01 ? shifted_values : values; + int l = nearest_int(0.5f*(id*scales[ib]+127.f)); + l = std::max(0, std::min(127, l)) << 1; + //printf("d = %g, id = %g, scales = %g, l = %d, dl = %g\n", d, id, scales[ib], l, d*(l - 127)); + y[ibl].scales[ib] |= l; + l -= 127; + float dl = d * l; + float idl = dl ? 1/dl : 0.f; + const float * xb = xbl + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ibl*super_block_size + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + auto qs = y[ibl].qs + ib*(block_size/2); + for (int j = 0; j < block_size/2; ++j) { + uint8_t i1 = best_index_iq4nl(block_values, idl*xb[j]); + uint8_t i2 = best_index_iq4nl(block_values, idl*xb[j+block_size/2]); + qs[j] = i1 | (i2 << 4); + float w1 = weight[j]; + float w2 = weight[j+block_size/2]; + float q1 = block_values[i1]*l; + float q2 = block_values[i2]*l; + sumqx += w1*q1*xb[j] + w2*q2*xb[j+block_size/2]; + sumq2 += w1*q1*q1 + w2*q2*q2; + //float diff = xb[j] - d*q1; mse += diff*diff; + //diff = xb[j+block_size/2] - d*q2; mse += diff*diff; + } + } + } + //printf("rmse = %g\n", sqrt(mse/n_per_row)); + if (sumq2 > 0) *dptr = sumqx/sumq2; +} +} + +void quantize_row_iq4_ks_ref(const float * x, block_iq4_ks * y, int64_t k) { + quantize_iq4_ks(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq4_ks(const float * x, void * y, int64_t k) { + quantize_iq4_ks(x, (void *)y, 1, k, nullptr, nullptr); +} + +size_t quantize_iq4_ks(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + constexpr int kBlockSize = 32; + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ4_KS, n_per_row); + float weight[kBlockSize]; + std::vector all_scales(n_per_row/kBlockSize); + QHelper helper(imatrix, user_data, n_per_row, kBlockSize); + auto q_func = [&all_scales, &weight, block_size = kBlockSize] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + quantize_row_iq4_k_impl_bs128(QK_K, block_size, n_per_row, x, (char *)vy, all_scales.data(), weight, iq4k_values, imatrix, 7); + }; + helper.quantize(nrows, src, dst, row_size, q_func); + return nrows * row_size; +} + +void dequantize_row_iq4_ks(const block_iq4_ks * x, float * y, int64_t k) { + constexpr int kBlockSize = 32; //128; + GGML_ASSERT(k%QK_K == 0); + const float * dptr = (const float *)x; + float d = *dptr; + x = (const block_iq4_ks *)(dptr + 1); + int nblock = k/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto qs = x[ibl].qs; + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + float dl = d * ((int)(x[ibl].scales[ib] & 254) - 127); + const int8_t * values = iq4k_values + ((x[ibl].scales[ib] & 1) << 4); + for (int j = 0; j < kBlockSize/2; ++j) { + y[j ] = dl * values[qs[j] & 0xf]; + y[j+kBlockSize/2] = dl * values[qs[j] >> 4]; + } + y += kBlockSize; + qs += kBlockSize/2; + } + } +} + +void vec_dot_iq4_ks_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + constexpr int kBlockSize = 32; +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ4_KS, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + const float * dptr = (const float *)vx; + const float d = *dptr; + //printf("%s: n = %d, d = %g\n", __func__, n, d); + const block_iq4_ks * x = (const block_iq4_ks *)(dptr + 1); + const block_q8_K * y = (const block_q8_K *)vy; + int nblock = n/QK_K; + float sumf = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + //int sumi = 0; + auto qy = y[ibl].qs; + auto qx = x[ibl].qs; + float db = d * y[ibl].d; + for (int ib = 0; ib < QK_K/kBlockSize; ++ib) { + float dl = db * ((x[ibl].scales[ib] & 254) - 127); + //int ls = (x[ibl].scales[ib] & 254) - 127; + const int8_t * values = iq4k_values + ((x[ibl].scales[ib] & 1) << 4); + int suml = 0; + for (int j = 0; j < kBlockSize/2; ++j) { + suml += qy[j ] * values[qx[j] & 0xf] + + qy[j + kBlockSize/2] * values[qx[j] >> 4]; + } + sumf += dl * suml; + //sumi += ls * suml; + qy += kBlockSize; + qx += kBlockSize/2; + } + //sumf += d * y[ibl].d * sumi; + } + *s = sumf; +} + +namespace { +static void quantize_row_iq5_ks_impl(const int super_block_size, const int block_size, + int n_per_row, const float * x, char * cy, + float * all_scales, float * weight, + const int8_t * values, + const float * quant_weights, + const int ntry) { + + float * dptr = (float *)cy; + dptr[0] = 0; + block_iq5_ks * y = (block_iq5_ks *)(dptr + 1); + + const int8_t * shifted_values = values + 32; + + float amax_scale = 0; + + for (int ibl = 0; ibl < n_per_row/super_block_size; ++ibl) { + memset(&y[ibl], 0, sizeof(block_iq5_ks)); + const float * xbl = x + ibl*super_block_size; + auto scales = all_scales + ibl*(super_block_size/block_size); + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= 2.f/super_block_size; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const float * xb = xbl + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ibl*super_block_size + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < block_size; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float d = ntry > 0 ? -max/values[0] : max/values[0]; + float id = 1/d; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq5nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq5nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + d = sumqx_p/sumq2_p; + bool is_shifted = false; + float best = d*sumqx_p; + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d*sumqx_m; + } + for (int itry = -ntry; itry <= ntry; ++itry) { + id = (itry + values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq5nl(values, al); + float q = values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq5nl(values, -al); + q = values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = false; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = false; + } + id = (itry + shifted_values[0])/max; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int j = 0; j < block_size; ++j) { + float w = weight[j]; + float al = id*xb[j]; + int l = best_index_iq5nl(shifted_values, al); + float q = shifted_values[l]; + sumqx_p += w*q*xb[j]; + sumq2_p += w*q*q; + l = best_index_iq5nl(shifted_values, -al); + q = shifted_values[l]; + sumqx_m += w*q*xb[j]; + sumq2_m += w*q*q; + } + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; is_shifted = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; is_shifted = true; + } + } + if (is_shifted) y[ibl].scales[ib] = 0x01; + scales[ib] = d; + amax_scale = std::max(amax_scale, std::abs(d)); + } + } + float d = amax_scale/127; + *dptr = d; + if (!d) return; + float id = d ? 1/d : 0.f; + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < n_per_row/super_block_size; ++ibl) { + const float * xbl = x + ibl*super_block_size; + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= 2.f/super_block_size; + auto scales = all_scales + (super_block_size/block_size)*ibl; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const int8_t * block_values = y[ibl].scales[ib] & 0x01 ? shifted_values : values; + int l = nearest_int(0.5f*(id*scales[ib]+127.f)); + l = std::max(0, std::min(127, l)) << 1; + y[ibl].scales[ib] |= l; + l -= 127; + float dl = d * l; + float idl = dl ? 1/dl : 0.f; + const float * xb = xbl + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ibl*super_block_size + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + for (int j = 0; j < block_size; ++j) { + uint8_t idx = best_index_iq5nl(block_values, idl*xb[j]); + y[ibl].qs[block_size*(ib/2) + j] |= ((idx & 0xf) << 4*(ib%2)); + y[ibl].qh[j] |= ((idx >> 4) << ib); + float w = weight[j]; + float q = block_values[idx]*l; + sumqx += w*q*xb[j]; + sumq2 += w*q*q; + } + } + } + if (sumq2 > 0) *dptr = sumqx/sumq2; +} +} + +void quantize_row_iq5_ks_ref(const float * x, block_iq5_ks * y, int64_t k) { + quantize_iq5_ks(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq5_ks(const float * x, void * y, int64_t k) { + quantize_iq5_ks(x, (void *)y, 1, k, nullptr, nullptr); +} + +size_t quantize_iq5_ks(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + constexpr int kBlockSize = 32; + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ5_KS, n_per_row); + float weight[kBlockSize]; + std::vector all_scales(n_per_row/kBlockSize); + QHelper helper(imatrix, user_data, n_per_row, kBlockSize); + auto q_func = [&all_scales, &weight, block_size = kBlockSize] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + quantize_row_iq5_ks_impl(QK_K, block_size, n_per_row, x, (char *)vy, all_scales.data(), weight, iq5nl_values, imatrix, 5); + }; + helper.quantize(nrows, src, dst, row_size, q_func); + return nrows * row_size; +} + +void dequantize_row_iq5_ks(const block_iq5_ks * x, float * y, int64_t k) { + constexpr int kBlockSize = 32; + GGML_ASSERT(k%QK_K == 0); + const float * dptr = (const float *)x; + float d = *dptr; + x = (const block_iq5_ks *)(dptr + 1); + int nblock = k/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto qs = x[ibl].qs; + auto qh = x[ibl].qh; + for (int ib64 = 0; ib64 < QK_K/(2*kBlockSize); ++ib64) { + float dl1 = d * ((int)(x[ibl].scales[2*ib64+0] & 254) - 127); + float dl2 = d * ((int)(x[ibl].scales[2*ib64+1] & 254) - 127); + const int8_t * values1 = iq5nl_values + ((x[ibl].scales[2*ib64+0] & 1) << 5); + const int8_t * values2 = iq5nl_values + ((x[ibl].scales[2*ib64+1] & 1) << 5); + for (int j = 0; j < kBlockSize; ++j) { + y[j ] = dl1 * values1[(qs[j] & 0xf) | (((qh[j] >> (2*ib64+0)) & 1) << 4)]; + y[j+kBlockSize] = dl2 * values2[(qs[j] >> 4) | (((qh[j] >> (2*ib64+1)) & 1) << 4)]; + } + y += 2*kBlockSize; + qs += kBlockSize; + } + } +} + +void vec_dot_iq5_ks_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + constexpr int kBlockSize = 32; +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ5_KS, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + const float * dptr = (const float *)vx; + const float d = *dptr; + const block_iq5_ks * x = (const block_iq5_ks *)(dptr + 1); + const block_q8_K * y = (const block_q8_K *)vy; + int nblock = n/QK_K; + float sumf = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto qy = y[ibl].qs; + auto qs = x[ibl].qs; + auto qh = x[ibl].qh; + float db = d * y[ibl].d; + for (int ib64 = 0; ib64 < QK_K/(2*kBlockSize); ++ib64) { + float dl1 = db * ((int)(x[ibl].scales[2*ib64+0] & 254) - 127); + float dl2 = db * ((int)(x[ibl].scales[2*ib64+1] & 254) - 127); + const int8_t * values1 = iq5nl_values + ((x[ibl].scales[2*ib64+0] & 1) << 5); + const int8_t * values2 = iq5nl_values + ((x[ibl].scales[2*ib64+1] & 1) << 5); + int suml1 = 0; + int suml2 = 0; + for (int j = 0; j < kBlockSize; ++j) { + suml1 += qy[j ] * values1[(qs[j] & 0xf) | (((qh[j] >> (2*ib64+0)) & 1) << 4)]; + suml2 += qy[j+kBlockSize] * values2[(qs[j] >> 4) | (((qh[j] >> (2*ib64+1)) & 1) << 4)]; + } + sumf += dl1*suml1 + dl2*suml2; + y += 2*kBlockSize; + qs += kBlockSize; + } + } + *s = sumf; +} + +namespace { +const uint16_t * scramble_table() { + static std::mutex mutex; + static std::vector table; + std::lock_guard lock(mutex); + if (table.empty()) { + table.resize(1 << 15); + for (int i = 0; i < int(table.size()); ++i) { + uint16_t val = i; + int non = popcount(val); + if (non%2) val |= (1 << 15); + bool found = false; + for (int j = 0; j < int(table.size()); ++j) { + if ((j ^ (j << 1)) == val) { + table[i] = j; found = true; break; + } + } + if (!found) { + printf("Oops: did not find for %d %u\n", i, val); + exit(1); + } + } + } + return table.data(); +} +uint16_t prune_iq4ks(uint16_t v, const int8_t * values, const float * x, const float * w, float dl) { + if (popcount(v)%2 == 0) return v; + float best_score = std::numeric_limits::max(); + uint8_t q4[4]; + int jbest = -1; + uint8_t bestq = 0; + for (int j = 0; j < 4; ++j) { + uint8_t q = (v >> 4*j) & 0xf; + q4[j] = q; + auto pc = popcount(q); + float diff0 = dl*iq4k_values[q] - x[j]; + int qmin = std::max(int(q)-2, 0); + int qmax = std::min(int(q)+2, 15); + for (int iq = qmin; iq <= qmax; ++iq) { + uint8_t qq = iq; + if (qq == q) continue; + int pci = popcount(qq); + if (std::abs(pci - pc)%2) { + float diff1 = dl*values[qq] - x[j]; + float score = w[j]*(diff1*diff1 - diff0*diff0); + if (score < best_score) { + best_score = score; jbest = j; bestq = qq; + } + } + } + } + GGML_ASSERT(jbest >= 0); + q4[jbest] = bestq; + return (q4[0] | (q4[1] << 4) | (q4[2] << 8) | (q4[3] << 12)); +} +static void quantize_row_iq4_kss_impl(int n_per_row, const float * x, char * cy, + float * all_scales, float * weight, + const int8_t * values, + const float * quant_weights, + const uint16_t * table, + const int ntry) { + + constexpr int super_block_size = 256; + constexpr int block_size = 32; + + float * dptr = (float *)cy; + *dptr = 0; + block_iq4_kss * y = (block_iq4_kss *)(dptr + 1); + + const int8_t * shifted_values = values + 16; + + uint16_t vps[block_size/2], vms[block_size/2], vs[block_size/2]; + float xv[4], wv[4]; + + float amax_scale = 0; + + for (int ibl = 0; ibl < n_per_row/super_block_size; ++ibl) { + memset(&y[ibl], 0, sizeof(block_iq4_kss)); + const float * xbl = x + ibl*super_block_size; + auto scales = all_scales + ibl*(super_block_size/block_size); + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= 2.f/super_block_size; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const float * xb = xbl + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ibl*super_block_size + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + float amax = 0, max = 0; + for (int j = 0; j < block_size; ++j) { + float ax = fabsf(xb[j]); + if (ax > amax) { + amax = ax; max = xb[j]; + } + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float best = 0; + float d = -max/iq4k_values[0]; + std::memset(vs, 0, block_size); + for (int itry = -ntry; itry <= ntry; ++itry) { + float id = (itry + values[0])/max; + float sumqx_p = 0, sumq2_p = 0; + float sumqx_m = 0, sumq2_m = 0; + float this_d = 1/id; + for (int k = 0; k < block_size/4; ++k) { + xv[0] = xb[2*k+0]; xv[1] = xb[2*k+0+block_size/2]; xv[2] = xb[2*k+1]; xv[3] = xb[2*k+1+block_size/2]; + wv[0] = weight[2*k+0]; wv[1] = weight[2*k+0+block_size/2]; wv[2] = weight[2*k+1]; wv[3] = weight[2*k+1+block_size/2]; + uint16_t vp = 0, vm = 0; + for (int j = 0; j < 4; ++j) { + float al = id*xv[j]; + vp |= (best_index_iq4nl(values, al) << 4*j); + vm |= (best_index_iq4nl(values, -al) << 4*j); + } + vp = prune_iq4ks(vp, values, xv, wv, this_d); + vm = prune_iq4ks(vm, values, xv, wv, this_d); + for (int j = 0; j < 4; ++j) { + float w = wv[j]; + float q = values[(vp >> 4*j) & 0xf]; + sumqx_p += w*q*xv[j]; + sumq2_p += w*q*q; + q = values[(vm >> 4*j) & 0xf]; + sumqx_m += w*q*xv[j]; + sumq2_m += w*q*q; + } + vps[k] = vp; + vms[k] = vm; + } + bool copy_p = false, copy_m = false; + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; copy_p = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; copy_m = true; + } + if (copy_m) { + std::memcpy(vs, vms, block_size); + } else if (copy_p) { + std::memcpy(vs, vps, block_size); + } + + id = (itry + shifted_values[0])/max; + this_d = 1/id; + sumqx_p = sumq2_p = 0; + sumqx_m = sumq2_m = 0; + for (int k = 0; k < block_size/4; ++k) { + xv[0] = xb[2*k+0]; xv[1] = xb[2*k+0+block_size/2]; xv[2] = xb[2*k+1]; xv[3] = xb[2*k+1+block_size/2]; + wv[0] = weight[2*k+0]; wv[1] = weight[2*k+0+block_size/2]; wv[2] = weight[2*k+1]; wv[3] = weight[2*k+1+block_size/2]; + uint16_t vp = 0, vm = 0; + for (int j = 0; j < 4; ++j) { + float al = id*xv[j]; + vp |= (best_index_iq4nl(shifted_values, al) << 4*j); + vm |= (best_index_iq4nl(shifted_values, -al) << 4*j); + } + vp = prune_iq4ks(vp, shifted_values, xv, wv, this_d); + vm = prune_iq4ks(vm, shifted_values, xv, wv, this_d); + for (int j = 0; j < 4; ++j) { + float w = wv[j]; + float q = shifted_values[(vp >> 4*j) & 0xf]; + sumqx_p += w*q*xv[j]; + sumq2_p += w*q*q; + q = shifted_values[(vm >> 4*j) & 0xf]; + sumqx_m += w*q*xv[j]; + sumq2_m += w*q*q; + } + vps[k] = vp; + vms[k] = vm; + } + copy_p = copy_m = false; + if (sumq2_p > 0 && sumqx_p*sumqx_p > best*sumq2_p) { + d = sumqx_p/sumq2_p; best = d * sumqx_p; copy_p = true; + } + if (sumq2_m > 0 && sumqx_m*sumqx_m > best*sumq2_m) { + d = sumqx_m/sumq2_m; best = d * sumqx_m; copy_m = true; + } + if (copy_m) { + std::memcpy(vs, vms, block_size); + } else if (copy_p) { + std::memcpy(vs, vps, block_size); + } + } + scales[ib] = d; + amax_scale = std::max(amax_scale, std::abs(d)); + } + } + float d = amax_scale/127; + *dptr = d; + if (!d) return; + float id = 1/d; + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < n_per_row/super_block_size; ++ibl) { + auto scales = all_scales + (super_block_size/block_size)*ibl; + const float * xbl = x + ibl*super_block_size; + float sigma2 = 0; + for (int j = 0; j < super_block_size; ++j) sigma2 += xbl[j]*xbl[j]; + sigma2 *= 2.f/super_block_size; + for (int ib = 0; ib < super_block_size/block_size; ++ib) { + const float * xb = xbl + ib*block_size; + if (quant_weights) { + const float * qw = quant_weights + ibl*super_block_size + ib*block_size; + for (int j = 0; j < block_size; ++j) weight[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + } else { + for (int j = 0; j < block_size; ++j) weight[j] = xb[j]*xb[j]; + } + int l = nearest_int(0.5f*(id*scales[ib]+127.f)); + l = (std::max(0, std::min(127, l)) << 1) - 127; + if (l) { + float dl = d*l; + float idl = 1/dl; + float mse_p = 0, mse_m = 0; + for (int k = 0; k < block_size/4; ++k) { + xv[0] = xb[2*k+0]; xv[1] = xb[2*k+0+block_size/2]; xv[2] = xb[2*k+1]; xv[3] = xb[2*k+1+block_size/2]; + wv[0] = weight[2*k+0]; wv[1] = weight[2*k+0+block_size/2]; wv[2] = weight[2*k+1]; wv[3] = weight[2*k+1+block_size/2]; + uint16_t vp = 0, vm = 0; + for (int j = 0; j < 4; ++j) { + float al = idl*xv[j]; + vp |= (best_index_iq4nl( values, al) << 4*j); + vm |= (best_index_iq4nl(shifted_values, al) << 4*j); + } + vp = prune_iq4ks(vp, values, xv, wv, dl); + vm = prune_iq4ks(vm, shifted_values, xv, wv, dl); + for (int j = 0; j < 4; ++j) { + float w = wv[j]; + float q = values[(vp >> 4*j) & 0xf]; + mse_p += w*(xv[j] - dl*q)*(xv[j] - dl*q); + q = shifted_values[(vm >> 4*j) & 0xf]; + mse_m += w*(xv[j] - dl*q)*(xv[j] - dl*q); + } + vps[k] = vp; + vms[k] = vm; + } + const uint16_t * v = vps; + const int8_t * block_values = values; + if (mse_m < mse_p) { + v = vms; + block_values = values + 16; + } + for (int k = 0; k < block_size/4; ++k) { + xv[0] = xb[2*k+0]; xv[1] = xb[2*k+0+block_size/2]; xv[2] = xb[2*k+1]; xv[3] = xb[2*k+1+block_size/2]; + wv[0] = weight[2*k+0]; wv[1] = weight[2*k+0+block_size/2]; wv[2] = weight[2*k+1]; wv[3] = weight[2*k+1+block_size/2]; + for (int j = 0; j < 4; ++j) { + float q = block_values[(v[k] >> 4*j) & 0xf] * l; + sumqx += wv[j]*q*xv[j]; + sumq2 += wv[j]*q*q; + } + } + l += 127; + if (mse_m < mse_p) l |= 1; + uint16_t * q16 = (uint16_t *)y[ibl].qs + (block_size/4)*ib; + for (int k = 0; k < block_size/4; ++k) { + auto val = table[v[k] & 0x7fff]; + q16[k] = (val << 1) | ((l >> k) & 1); + } + } else { + l += 127; + uint16_t * q16 = (uint16_t *)y[ibl].qs + (block_size/4)*ib; + for (int k = 0; k < block_size/4; ++k) { + q16[k] = ((l >> k) & 1); + } + } + } + } + if (sumq2 > 0) *dptr = sumqx/sumq2 * 1.01f; +} +} + +size_t quantize_iq4_kss(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + constexpr int kBlockSize = 32; + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ4_KSS, n_per_row); + std::vector all_scales(n_per_row/kBlockSize); + float weight[kBlockSize]; + auto table = scramble_table(); + QHelper helper(imatrix, user_data, n_per_row, kBlockSize); + auto q_func = [&all_scales, &weight, table] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + quantize_row_iq4_kss_impl(n_per_row, x, (char *)vy, all_scales.data(), weight, iq4k_values, imatrix, table, 7); + }; + helper.quantize(nrows, src, dst, row_size, q_func); + return nrows * row_size; +} + +void quantize_row_iq4_kss_ref(const float * x, block_iq4_kss * y, int64_t k) { + quantize_iq4_kss(x, y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq4_kss(const float * x, void * y, int64_t k) { + quantize_iq4_kss(x, (block_iq4_kss *)y, 1, k, nullptr, nullptr); +} + +void dequantize_row_iq4_kss(const block_iq4_kss * x, float * y, int64_t k) { + const float * dptr = (const float *)x; + const float d = *dptr; + x = (const block_iq4_kss *)(dptr + 1); + uint16_t aux16[8]; + const uint8_t * aux8 = (const uint8_t *)aux16; + for (int ibl = 0; ibl < k/QK_K; ++ibl) { + auto qs = (const uint16_t *)x[ibl].qs; + for (int ib = 0; ib < QK_K/32; ++ib) { + int16_t ls = 0; + for (int k = 0; k < 8; ++k) { + aux16[k] = qs[k] & 0xfffe; + aux16[k] ^= (aux16[k] >> 1); + ls |= (qs[k] & 1) << k; + } + const int8_t * values = iq4k_values + ((ls & 1) << 4); + float dl = d * ((ls & 254) - 127); + for (int j = 0; j < 16; ++j) { + y[j+ 0] = dl * values[aux8[j] & 0xf]; + y[j+16] = dl * values[aux8[j] >> 4]; + } + y += 32; + qs += 8; + } + } +} + +void vec_dot_iq4_kss_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ4_KSS, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK_K == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq4_nl_r4 +// +void quantize_row_iq4_nl_r4_ref(const float * x, block_iq4_nl_r4 * y, int64_t k) { + // we assume we are called with 4 rows + quantize_iq4_nl_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq4_nl_r4(const float * x, void * y, int64_t k) { + // we assume we are called with 4 rows + quantize_iq4_nl_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static void repack_iq4_nl(int nrows, int n_per_row, const block_iq4_nl * x, block_iq4_nl_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK4_NL == 0); + int nblock = n_per_row/QK4_NL; + const block_iq4_nl * x4[4]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ib = 0; ib < nblock; ++ib) { + for (int k = 0; k < 4; ++k) y[ib].d[k] = x4[k][ib].d; + for (int k = 0; k < 4; ++k) for (int i = 0; i < 4; ++i) { + y[ib].qs[4*k+i+ 0] = (x4[k][ib].qs[i+0] & 0xf) | ((x4[k][ib].qs[i+ 8] & 0x0f) << 4); // 0....3 + 8...11 from each row + y[ib].qs[4*k+i+16] = (x4[k][ib].qs[i+0] >> 4) | ((x4[k][ib].qs[i+ 8] & 0xf0)); // 16...19 + 24...27 from each row + y[ib].qs[4*k+i+32] = (x4[k][ib].qs[i+4] & 0xf) | ((x4[k][ib].qs[i+12] & 0x0f) << 4); // 4....7 + 12...15 from each row + y[ib].qs[4*k+i+48] = (x4[k][ib].qs[i+4] >> 4) | ((x4[k][ib].qs[i+12] & 0xf0)); // 20...23 + 28...31 from each row + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq4_nl_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + auto row_size_nl = ggml_row_size(GGML_TYPE_IQ4_NL, n_per_row); + std::vector qtmp(4*row_size_nl); + QHelper helper(imatrix, user_data, n_per_row, 32); + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_iq4_nl(x, (char *)vy, 1, n_per_row, imatrix, nullptr); + }; + char * qrow = (char *)dst; + for (int row = 0; row < nrows; row += 4) { + helper.quantize(4, src, qtmp.data(), row_size_nl, q_func); + repack_iq4_nl(4, n_per_row, (const block_iq4_nl *)qtmp.data(), (block_iq4_nl_r4 *)qrow, false); + src += 4*n_per_row; + qrow += 4*row_size_nl; + } + return nrows*row_size_nl; +} + +void dequantize_row_iq4_nl_r4(const block_iq4_nl_r4 * x, float * y, int64_t k) { + // we assume we are called with 4 rows + int n_per_row = k/4; + int nb = n_per_row/QK4_NL; + float * yk[4]; + for (int k = 0; k < 4; ++k) yk[k] = y + k*n_per_row; + for (int ib = 0; ib < nb; ++ib) { + for (int k = 0; k < 4; ++k) { + float scale = GGML_FP16_TO_FP32(x[ib].d[k]); + for (int i = 0; i < 4; ++i) { + yk[k][QK4_NL*ib+i+ 0] = scale * iq4k_values[x[ib].qs[4*k+i+ 0] & 0xf]; + yk[k][QK4_NL*ib+i+ 8] = scale * iq4k_values[x[ib].qs[4*k+i+ 0] >> 4]; + yk[k][QK4_NL*ib+i+16] = scale * iq4k_values[x[ib].qs[4*k+i+16] & 0xf]; + yk[k][QK4_NL*ib+i+24] = scale * iq4k_values[x[ib].qs[4*k+i+16] >> 4]; + yk[k][QK4_NL*ib+i+ 4] = scale * iq4k_values[x[ib].qs[4*k+i+32] & 0xf]; + yk[k][QK4_NL*ib+i+12] = scale * iq4k_values[x[ib].qs[4*k+i+32] >> 4]; + yk[k][QK4_NL*ib+i+20] = scale * iq4k_values[x[ib].qs[4*k+i+48] & 0xf]; + yk[k][QK4_NL*ib+i+28] = scale * iq4k_values[x[ib].qs[4*k+i+48] >> 4]; + } + } + } +} + +void vec_dot_iq4_nl_r4_q8_0(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ4_NL_R4, vx, 0, GGML_TYPE_Q8_0, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q4_0_r8 +// +void quantize_row_q4_0_r8_ref(const float * x, block_iq4_nl_r8 * y, int64_t k) { + // we assume we are called with 8 rows + quantize_q4_0_r8(x, (void *)y, 8, k/8, nullptr, nullptr); +} + +void quantize_row_q4_0_r8(const float * x, void * y, int64_t k) { + // we assume we are called with 8 rows + quantize_q4_0_r8(x, y, 8, k/8, nullptr, nullptr); +} + +static void repack_q4_0(int nrows, int n_per_row, const block_q4_0 * x, block_iq4_nl_r8 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%8 == 0); + GGML_ASSERT(n_per_row%QK4_0 == 0); + int nblock = n_per_row/QK4_0; + const block_q4_0 * x8[8]; + for (int row = 0; row < nrows; row += 8) { + for (int k = 0; k < 8; ++k) x8[k] = x + nblock*k; + for (int ib = 0; ib < nblock; ++ib) { + for (int k = 0; k < 8; ++k) { + y[ib].d[k] = x8[k][ib].d; + for (int l = 0; l < 4; ++l) { + for (int i = 0; i < 4; ++i) { + y[ib].qs[32*l+4*k+i] = x8[k][ib].qs[4*l + i]; + } + } + } +#ifdef __ARM_NEON + if (online) { + for (int l = 0; l < 8; ++l) { + auto v = vld1q_u8(y[ib].qs + 16*l); + vst1q_u8(y[ib].qs + 16*l, veorq_u8(v, vdupq_n_u8(0x88))); + } + } +#endif + } + x += 8*nblock; + y += nblock; + } +} +#ifdef __ARM_NEON +static void modify_q4_0_r8(int64_t k, char * cy) { + auto y = (block_iq4_nl_r8 *)cy; + int nb = k/(32*8); + for (int ib = 0; ib < nb; ++ib) { + auto v1 = vld1q_u8_x4(y[ib].qs); + auto v2 = vld1q_u8_x4(y[ib].qs+64); + for (int j = 0; j < 4; ++j) { + v1.val[j] = veorq_u8(v1.val[j], vdupq_n_u8(0x88)); + v2.val[j] = veorq_u8(v2.val[j], vdupq_n_u8(0x88)); + } + vst1q_u8_x4(y[ib].qs+ 0, v1); + vst1q_u8_x4(y[ib].qs+64, v2); + } +} +#endif + +size_t quantize_q4_0_r8(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + GGML_ASSERT(nrows%8 == 0); + auto row_size_nl = ggml_row_size(GGML_TYPE_Q4_0, n_per_row); + std::vector qtmp(8*row_size_nl); + QHelper helper(imatrix, user_data, n_per_row, 32); + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_q4_0(x, (char *)vy, 1, n_per_row, imatrix, nullptr); + }; + char * qrow = (char *)dst; + for (int row = 0; row < nrows; row += 8) { + helper.quantize(8, src, qtmp.data(), row_size_nl, q_func); + repack_q4_0(8, n_per_row, (const block_q4_0 *)qtmp.data(), (block_iq4_nl_r8 *)qrow, false); + src += 8*n_per_row; + qrow += 8*row_size_nl; + } + return nrows*row_size_nl; +} + +void dequantize_row_q4_0_r8(const block_iq4_nl_r8 * x, float * y, int64_t k) { + // we assume we are called with 8 rows + int n_per_row = k/8; + int nb = n_per_row/QK4_0; + float * yk[8]; + for (int k = 0; k < 8; ++k) yk[k] = y + k*n_per_row; + for (int ib = 0; ib < nb; ++ib) { + for (int k = 0; k < 8; ++k) { + float scale = GGML_FP16_TO_FP32(x[ib].d[k]); + for (int l = 0; l < 4; ++l) { + for (int i = 0; i < 4; ++i) { + yk[k][QK4_0*ib+4*l+i+ 0] = scale * ((x[ib].qs[32*l+4*k+i] & 0xf) - 8); + yk[k][QK4_0*ib+4*l+i+16] = scale * ((x[ib].qs[32*l+4*k+i] >> 4) - 8); + } + } + } + } +} + +void vec_dot_q4_0_r8_q8_0(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q4_0_R8, vx, 0, GGML_TYPE_Q8_0, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + + +// +// ========================================= q8_0_r8 +// +void quantize_row_q8_0_r8_ref(const float * x, block_q8_0_r8 * y, int64_t k) { + // we assume we are called with 4 rows + quantize_q8_0_r8(x, (void *)y, 8, k/8, nullptr, nullptr); +} + +void quantize_row_q8_0_r8(const float * x, void * y, int64_t k) { + // we assume we are called with 4 rows + quantize_q8_0_r8(x, y, 8, k/8, nullptr, nullptr); +} + +static void repack_q8_0(int nrows, int n_per_row, const block_q8_0 * x, block_q8_0_r8 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%8 == 0); + GGML_ASSERT(n_per_row%QK8_0 == 0); + int nblock = n_per_row/QK8_0; + const block_q8_0 * x8[8]; + for (int row = 0; row < nrows; row += 8) { + for (int k = 0; k < 8; ++k) x8[k] = x + nblock*k; + for (int ib = 0; ib < nblock; ++ib) { + for (int k = 0; k < 8; ++k) y[ib].d[k] = x8[k][ib].d; + for (int l = 0; l < 4; ++l) { + for (int k = 0; k < 8; ++k) for (int i = 0; i < 4; ++i) { + y[ib].qs[32*l+4*k+i+ 0] = x8[k][ib].qs[i+4*l+ 0]; + y[ib].qs[32*l+4*k+i+128] = x8[k][ib].qs[i+4*l+16]; + } + } +#ifdef HAVE_FANCY_SIMD + if (online) { + for (int l = 0; l < 4; ++l) { + auto v = _mm512_add_epi8(_mm512_loadu_si512((const __m512i *)y[ib].qs + l), _mm512_set1_epi8(127)); + _mm512_storeu_si512((__m512i *)y[ib].qs + l, v); + } + } +#endif + } + x += 8*nblock; + y += nblock; + } +} + +#ifdef HAVE_FANCY_SIMD +static void modify_q8_0_r8(int64_t k, char * cy) { + auto y = (block_q8_0_r8 *)cy; + int nb = k/(32*8); + for (int ib = 0; ib < nb; ++ib) { + for (int l = 0; l < 4; ++l) { + auto v = _mm512_add_epi8(_mm512_loadu_si512((const __m512i *)y[ib].qs + l), _mm512_set1_epi8(127)); + _mm512_storeu_si512((__m512i *)y[ib].qs + l, v); + } + } +} +#endif + +size_t quantize_q8_0_r8(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%8 == 0); + auto row_size_0 = ggml_row_size(GGML_TYPE_Q8_0, n_per_row); + std::vector qtmp(8*row_size_0); + char * qrow = (char *)dst; + for (int row = 0; row < nrows; row += 8) { + iqkbase_quantize_q8_0(src, qtmp.data(), 8, n_per_row, imatrix, nullptr); + repack_q8_0(8, n_per_row, (const block_q8_0 *)qtmp.data(), (block_q8_0_r8 *)qrow, false); + src += 8*n_per_row; + qrow += 8*row_size_0; + } + return nrows*row_size_0; +} + +void dequantize_row_q8_0_r8(const block_q8_0_r8 * x, float * y, int64_t k) { + // we assume we are called with 4 rows + int n_per_row = k/8; + int nb = n_per_row/QK8_0; + float * yk[8]; + for (int k = 0; k < 8; ++k) yk[k] = y + k*n_per_row; + for (int ib = 0; ib < nb; ++ib) { + for (int k = 0; k < 8; ++k) { + float scale = GGML_FP16_TO_FP32(x[ib].d[k]); + for (int l = 0; l < 4; ++l) for (int i = 0; i < 4; ++i) { + yk[k][QK8_0*ib+4*l+i+ 0] = scale * x[ib].qs[32*l+4*k+i+ 0]; + yk[k][QK8_0*ib+4*l+i+16] = scale * x[ib].qs[32*l+4*k+i+128]; + } + } + } +} + +void vec_dot_q8_0_r8_q8_0(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q8_0_R8, vx, 0, GGML_TYPE_Q8_0, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q5_0_r4 +// +void quantize_row_q5_0_r4_ref(const float * x, block_q5_0_r4 * y, int64_t k) { + // we assume we are called with 4 rows + quantize_q5_0_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_q5_0_r4(const float * x, void * y, int64_t k) { + // we assume we are called with 4 rows + quantize_q5_0_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static inline void convert_q5_0(const block_q5_0& x, uint8_t * L) { + uint32_t qh; + memcpy(&qh, x.qh, sizeof(qh)); + + for (int j = 0; j < QK5_0/2; ++j) { + const uint8_t xh_0 = ((qh >> (j + 0)) << 4) & 0x10; + const uint8_t xh_1 = ((qh >> (j + 12)) ) & 0x10; + + L[j ] = (x.qs[j] & 0x0F) | xh_0; + L[j + QK4_0/2] = (x.qs[j] >> 4) | xh_1; + } +} + +static void repack_q5_0(int nrows, int n_per_row, const block_q5_0 * x, block_q5_0_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK5_0 == 0); + int nblock = n_per_row/QK5_0; + const block_q5_0 * x4[4]; + uint8_t L[QK5_0]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ib = 0; ib < nblock; ++ib) { + std::memset(y[ib].qh, 0, QK5_0/2); + for (int k = 0; k < 4; ++k) { + y[ib].d[k] = x4[k][ib].d; + convert_q5_0(x4[k][ib], L); + for (int l = 0; l < 4; ++l) { + int l1 = 4*(l/2) + 16*(l%2), l2 = l1 + 8; + for (int i = 0; i < 4; ++i) { + y[ib].qs[4*k+i+16*l] = (L[i + l1] & 0xf) | ((L[i + l2] & 0xf) << 4); + y[ib].qh[4*k+i] |= ((L[i + l1] >> 4) | ((L[i + l2] >> 4) << 4)) << l; + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_q5_0_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + auto row_size_0 = ggml_row_size(GGML_TYPE_Q5_0, n_per_row); + std::vector qtmp(4*row_size_0); + char * qrow = (char *)dst; + for (int row = 0; row < nrows; row += 4) { + iqkbase_quantize_q5_0(src, qtmp.data(), 4, n_per_row, imatrix, user_data); + repack_q5_0(4, n_per_row, (const block_q5_0 *)qtmp.data(), (block_q5_0_r4 *)qrow, false); + src += 4*n_per_row; + qrow += 4*row_size_0; + } + return nrows*row_size_0; +} + +void dequantize_row_q5_0_r4(const block_q5_0_r4 * x, float * y, int64_t k) { + // we assume we are called with 4 rows + int n_per_row = k/4; + int nb = n_per_row/QK8_0; + float * yk[4]; + for (int k = 0; k < 4; ++k) yk[k] = y + k*n_per_row; + for (int ib = 0; ib < nb; ++ib) { + for (int k = 0; k < 4; ++k) { + float d = GGML_FP16_TO_FP32(x[ib].d[k]); + float m = -16*d; + for (int l = 0; l < 4; ++l) { + int ll = 16*(l%2) + 4*(l/2); + for (int i = 0; i < 4; ++i) { + yk[k][QK4_0*ib+i+ll+0] = d * ((x[ib].qs[4*k+i+16*l] & 0xf) | (((x[ib].qh[4*k+i] >> (l+0)) & 1) << 4)) + m; + yk[k][QK4_0*ib+i+ll+8] = d * ((x[ib].qs[4*k+i+16*l] >> 4) | (((x[ib].qh[4*k+i] >> (l+4)) & 1) << 4)) + m; + } + } + } + } +} + +void vec_dot_q5_0_r4_q8_0(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q5_0_R4, vx, 0, GGML_TYPE_Q8_0, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q6_0_r4 +// +void quantize_row_q6_0_r4_ref(const float * x, block_q6_0_r4 * y, int64_t k) { + // we assume we are called with 4 rows + quantize_q6_0_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_q6_0_r4(const float * x, void * y, int64_t k) { + // we assume we are called with 4 rows + quantize_q6_0_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static inline void convert_q6_0(const block_q6_0& x, uint8_t * L) { + + for (int j = 0; j < QK6_0/2; ++j) { + const uint8_t h = x.qh[j%(QK6_0/4)] >> 4*(j/(QK6_0/4)); + L[j ] = (x.qs[j] & 0x0F) | ((h << 4) & 0x30); + L[j + QK6_0/2] = (x.qs[j] >> 4) | ((h << 2) & 0x30); + } +} + +static void repack_q6_0(int nrows, int n_per_row, const block_q6_0 * x, block_q6_0_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK5_0 == 0); + int nblock = n_per_row/QK6_0; + const block_q6_0 * x4[4]; + uint8_t L[QK6_0]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ib = 0; ib < nblock; ++ib) { + std::memset(y[ib].qh, 0, QK6_0); + for (int k = 0; k < 4; ++k) { + y[ib].d[k] = x4[k][ib].d; + convert_q6_0(x4[k][ib], L); + for (int l = 0; l < 4; ++l) { + int l1 = 4*(l/2) + 16*(l%2), l2 = l1 + 8; + for (int i = 0; i < 4; ++i) { + y[ib].qs[4*k+i+16*l] = (L[i + l1] & 0xf) | ((L[i + l2] & 0xf) << 4); + y[ib].qh[4*k+i+16*(l%2)] |= ((L[i + l1] >> 4) | ((L[i + l2] >> 4) << 4)) << 2*(l/2); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_q6_0_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + auto row_size_0 = ggml_row_size(GGML_TYPE_Q6_0, n_per_row); + std::vector qtmp(4*row_size_0); + char * qrow = (char *)dst; + QHelper helper(imatrix, user_data, n_per_row, 32); + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_q6_0(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + for (int row = 0; row < nrows; row += 4) { + helper.quantize(4, src, qtmp.data(), row_size_0, q_func); + repack_q6_0(4, n_per_row, (const block_q6_0 *)qtmp.data(), (block_q6_0_r4 *)qrow, false); + src += 4*n_per_row; + qrow += 4*row_size_0; + } + return nrows*row_size_0; +} + +void dequantize_row_q6_0_r4(const block_q6_0_r4 * x, float * y, int64_t k) { + // we assume we are called with 4 rows + int n_per_row = k/4; + int nb = n_per_row/QK6_0; + float * yk[4]; + for (int k = 0; k < 4; ++k) yk[k] = y + k*n_per_row; + for (int ib = 0; ib < nb; ++ib) { + for (int k = 0; k < 4; ++k) { + float d = GGML_FP16_TO_FP32(x[ib].d[k]); + float m = -32*d; + for (int l = 0; l < 4; ++l) { + int ll = 16*(l%2) + 4*(l/2); + for (int i = 0; i < 4; ++i) { + yk[k][QK4_0*ib+i+ll+0] = d * ((x[ib].qs[4*k+i+16*l] & 0xf) | (((x[ib].qh[4*k+i+16*(l%2)] >> (2*(l/2)+0)) & 3) << 4)) + m; + yk[k][QK4_0*ib+i+ll+8] = d * ((x[ib].qs[4*k+i+16*l] >> 4) | (((x[ib].qh[4*k+i+16*(l%2)] >> (2*(l/2)+4)) & 3) << 4)) + m; + } + } + } + } +} + +void vec_dot_q6_0_r4_q8_0(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q6_0_R4, vx, 0, GGML_TYPE_Q8_0, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq4_xs_r8 +// + +void quantize_row_iq4_xs_r8_ref(const float * x, block_iq4_xs_r8 * y, int64_t k) { + quantize_iq4_xs_r8(x, (void *)y, 8, k/8, nullptr, nullptr); +} + +void quantize_row_iq4_xs_r8(const float * x, void * y, int64_t k) { + quantize_iq4_xs_r8(x, y, 8, k/8, nullptr, nullptr); +} + +static void repack_iq4_xs(int nrows, int n_per_row, const block_iq4_xs * x, block_iq4_xs_r8 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%8 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq4_xs * x8[8]; + for (int row = 0; row < nrows; row += 8) { + for (int k = 0; k < 8; ++k) x8[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].scales_l, 0, QK_K/8); + std::memset(y[ibl].scales_h, 0, QK_K/16); + for (int k = 0; k < 8; ++k) { + y[ibl].d[k] = x8[k][ibl].d; + for (int ib = 0; ib < QK_K/32; ++ib) { + uint8_t sl = (x8[k][ibl].scales_l[ib/2] >> 4*(ib%2)) & 0xf; + uint8_t sh = (x8[k][ibl].scales_h >> 2*ib) & 3; + int i = 8*ib + k; + y[ibl].scales_l[i%32] |= (sl << 4*(i/32)); + y[ibl].scales_h[i%16] |= (sh << 2*(i/16)); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[128*ib+4*k+i+ 0] = (x8[k][ibl].qs[16*ib+i+0] & 0xf) | ((x8[k][ibl].qs[16*ib+i+ 4] & 0xf) << 4); + y[ibl].qs[128*ib+4*k+i+32] = (x8[k][ibl].qs[16*ib+i+8] & 0xf) | ((x8[k][ibl].qs[16*ib+i+12] & 0xf) << 4); + y[ibl].qs[128*ib+4*k+i+64] = (x8[k][ibl].qs[16*ib+i+0] >> 4) | ((x8[k][ibl].qs[16*ib+i+ 4] >> 4) << 4); + y[ibl].qs[128*ib+4*k+i+96] = (x8[k][ibl].qs[16*ib+i+8] >> 4) | ((x8[k][ibl].qs[16*ib+i+12] >> 4) << 4); + } + } + } + } + x += 8*nblock; + y += nblock; + } +} + +size_t quantize_iq4_xs_r8(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_iq4_xs(x, (char *)vy, 1, n_per_row, imatrix, nullptr); + }; + return quantize_repack<32, block_iq4_xs, block_iq4_xs_r8, 8>(GGML_TYPE_IQ4_XS, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_iq4_xs); +} + +void dequantize_row_iq4_xs_r8(const block_iq4_xs_r8 * x, float * y, int64_t k) { + auto n_per_row = k/8; + float * y8[8]; + for (int k = 0; k < 8; ++k) y8[k] = y + n_per_row*k; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 8; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 8*ib + k; + float dl = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 3) << 4)) - 32); + for (int l = 0; l < 4; ++l) for (int i = 0; i < 4; ++i) { + y8[k][QK_K*ibl+32*ib+8*l+i+0] = dl * iq4k_values[x[ibl].qs[128*ib+4*k+i+32*l] & 0xf]; + y8[k][QK_K*ibl+32*ib+8*l+i+4] = dl * iq4k_values[x[ibl].qs[128*ib+4*k+i+32*l] >> 4]; + } + } + } + } +} + +void vec_dot_iq4_xs_r8_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ4_XS_R8, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq4_ks_r4 +// + +void quantize_row_iq4_ks_r4_ref(const float * x, block_iq4_ks_r4 * y, int64_t k) { + quantize_iq4_ks_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq4_ks_r4(const float * x, void * y, int64_t k) { + quantize_iq4_ks_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static void repack_iq4_ks(int nrows, int n_per_row, const block_iq4_ks * x, block_iq4_ks_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ4_KS, n_per_row); + int nblock = n_per_row/QK_K; + char * cy = (char *)y; + const char * cx = (const char *)x; + const block_iq4_ks * x4[4]; + for (int row = 0; row < nrows; row += 4) { + float * dptr = (float *)cy; + block_iq4_ks_r4 * y = (block_iq4_ks_r4 *)(dptr + 4); + for (int k = 0; k < 4; ++k) { + auto dk = (const float *)(cx + k*row_size); + dptr[k] = dk[0]; + x4[k] = (const block_iq4_ks *)(dk + 1); + } + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + for (int ib = 0; ib < QK_K/32; ++ib) { + y[ibl].scales[4*ib+k] = x4[k][ibl].scales[ib]; + for (int i = 0; i < 4; ++i) { + y[ibl].qs[64*ib+4*k+i+ 0] = (x4[k][ibl].qs[16*ib+i+0] & 0xf) | ((x4[k][ibl].qs[16*ib+i+ 8] & 0x0f) << 4); // 0....3 + 8...11 from each row + y[ibl].qs[64*ib+4*k+i+16] = (x4[k][ibl].qs[16*ib+i+0] >> 4) | ((x4[k][ibl].qs[16*ib+i+ 8] & 0xf0)); // 16...19 + 24...27 from each row + y[ibl].qs[64*ib+4*k+i+32] = (x4[k][ibl].qs[16*ib+i+4] & 0xf) | ((x4[k][ibl].qs[16*ib+i+12] & 0x0f) << 4); // 4....7 + 12...15 from each row + y[ibl].qs[64*ib+4*k+i+48] = (x4[k][ibl].qs[16*ib+i+4] >> 4) | ((x4[k][ibl].qs[16*ib+i+12] & 0xf0)); // 20...23 + 28...31 from each row + } + } + } + } + cx += 4*row_size; + cy += 4*row_size; + } +} + +size_t quantize_iq4_ks_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + char * qcur = (char *)dst; + auto row_size = ggml_row_size(GGML_TYPE_IQ4_KS, n_per_row); + std::vector qtmp(4*row_size); + for (int row = 0; row < nrows; row += 4) { + quantize_iq4_ks(src, (void *)qtmp.data(), 4, n_per_row, imatrix, user_data); + repack_iq4_ks(4, n_per_row, (const block_iq4_ks *)qtmp.data(), (block_iq4_ks_r4 *)qcur, false); + qcur += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq4_ks_r4(const block_iq4_ks_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + const float * dptr = (const float *)x; + x = (const block_iq4_ks_r4 *)(dptr + 4); + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = dptr[k]; + for (int ib = 0; ib < QK_K/32; ++ib) { + float dl = d * ((x[ibl].scales[4*ib + k] & 254) - 127); + auto values = iq4k_values + ((x[ibl].scales[4*ib + k] & 1) << 4); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl * values[x[ibl].qs[64*ib+4*k+i+ 0] & 0xf]; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl * values[x[ibl].qs[64*ib+4*k+i+ 0] >> 4]; + y4[k][QK_K*ibl+32*ib+i+16] = dl * values[x[ibl].qs[64*ib+4*k+i+16] & 0xf]; + y4[k][QK_K*ibl+32*ib+i+24] = dl * values[x[ibl].qs[64*ib+4*k+i+16] >> 4]; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl * values[x[ibl].qs[64*ib+4*k+i+32] & 0xf]; + y4[k][QK_K*ibl+32*ib+i+12] = dl * values[x[ibl].qs[64*ib+4*k+i+32] >> 4]; + y4[k][QK_K*ibl+32*ib+i+20] = dl * values[x[ibl].qs[64*ib+4*k+i+48] & 0xf]; + y4[k][QK_K*ibl+32*ib+i+28] = dl * values[x[ibl].qs[64*ib+4*k+i+48] >> 4]; + } + } + } + } +} + +void vec_dot_iq4_ks_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ4_KS_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq2_bn_r4 +// +void quantize_row_iq2_bn_r4_ref(const float * x, block_iq2_bn * y, int64_t k) { + quantize_iq2_bn_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq2_bn_r4(const float * x, void * y, int64_t k) { + quantize_iq2_bn_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +void repack_iq2_bn(int nrows, int n_per_row, const char * x, char * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_IQ1BN == 0); + int nblock = n_per_row/QK_IQ1BN; + auto row_size = ggml_row_size(GGML_TYPE_IQ2_BN, n_per_row); + const uint8_t * x4[4]; + for (int row = 0; row < nrows; row += 4) { + float * dr4 = (float *)(y + 4*row*row_size); + for (int k = 0; k < 4; ++k) { + const float * dptr = (const float *)(x + (row + k)*row_size); + dr4[k] = *dptr; + x4[k] = (const uint8_t *)(dptr + 1); + } + uint8_t * y4 = (uint8_t *)(dr4 + 4); + //std::memset(y4, 0, n_per_row); + for (int ib = 0; ib < nblock; ++ib) { + // 0...3 from rows 0...3 go to 1st 2 bits of 0...15 + // 16..19 from rows 0...3 go to 1st 2 bits of 16...31 + // 32..35 from rows 0...3 go to 1st 2 bits of 32...47 + // 48..51 from rows 0...3 go to 1st 2 bits of 48...63 + // 4...7 from rows 0...3 go to 2nd 2 bits of 0...15 + // 20..23 from rows 0...3 go to 2nd 2 bits of 16...31 + // 36..39 from rows 0...3 go to 2nd 2 bits of 32...47 + // 52..55 from rows 0...3 go to 2nd 2 bits of 48...63 + // 8..11 from rows 0...3 go to 3rd 2 bits of 0...15 + // 24..27 from rows 0...3 go to 3rd 2 bits of 16...31 + // 40..43 from rows 0...3 go to 3rd 2 bits of 32...47 + // 56..59 from rows 0...3 go to 3rd 2 bits of 48...63 + // 12..15 from rows 0...3 go to 4th 2 bits of 0...15 + // 28..31 from rows 0...3 go to 4th 2 bits of 16...31 + // 44..47 from rows 0...3 go to 4th 2 bits of 32...47 + // 60..63 from rows 0...3 go to 4th 2 bits of 48...63 + for (int k = 0; k < 4; ++k) { + for (int l = 0; l < 4; ++l) for (int i = 0; i < 4; ++i) { + y4[64*ib + 4*k + i + 16*l] = (((x4[k][16*ib + i + 0] >> 2*l) & 3) << 0) | + (((x4[k][16*ib + i + 4] >> 2*l) & 3) << 2) | + (((x4[k][16*ib + i + 8] >> 2*l) & 3) << 4) | + (((x4[k][16*ib + i + 12] >> 2*l) & 3) << 6); + //y4[64*ib + 4*k + i + 0] |= (x4[k][16*ib + i] >> 0) & 3; + //y4[64*ib + 4*k + i + 16] |= (x4[k][16*ib + i] >> 2) & 3; + //y4[64*ib + 4*k + i + 32] |= (x4[k][16*ib + i] >> 4) & 3; + //y4[64*ib + 4*k + i + 48] |= (x4[k][16*ib + i] >> 6) & 3; + //y4[64*ib + 4*k + i + 0] |= ((x4[k][16*ib + i + 4] >> 0) & 3) << 2; + //y4[64*ib + 4*k + i + 16] |= ((x4[k][16*ib + i + 4] >> 2) & 3) << 2; + //y4[64*ib + 4*k + i + 32] |= ((x4[k][16*ib + i + 4] >> 4) & 3) << 2; + //y4[64*ib + 4*k + i + 48] |= ((x4[k][16*ib + i + 4] >> 6) & 3) << 2; + //y4[64*ib + 4*k + i + 0] |= ((x4[k][16*ib + i + 8] >> 0) & 3) << 4; + //y4[64*ib + 4*k + i + 16] |= ((x4[k][16*ib + i + 8] >> 2) & 3) << 4; + //y4[64*ib + 4*k + i + 32] |= ((x4[k][16*ib + i + 8] >> 4) & 3) << 4; + //y4[64*ib + 4*k + i + 48] |= ((x4[k][16*ib + i + 8] >> 6) & 3) << 4; + //y4[64*ib + 4*k + i + 0] |= ((x4[k][16*ib + i + 12] >> 0) & 3) << 6; + //y4[64*ib + 4*k + i + 16] |= ((x4[k][16*ib + i + 12] >> 2) & 3) << 6; + //y4[64*ib + 4*k + i + 32] |= ((x4[k][16*ib + i + 12] >> 4) & 3) << 6; + //y4[64*ib + 4*k + i + 48] |= ((x4[k][16*ib + i + 12] >> 6) & 3) << 6; + } + } + } + } +} +} + +size_t quantize_iq2_bn_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_IQ1BN == 0); + char * qcur = (char *)dst; + auto row_size = ggml_row_size(GGML_TYPE_IQ2_BN, n_per_row); + std::vector qtmp(4*row_size); + for (int row = 0; row < nrows; row += 4) { + quantize_iq2_bn(src, (void *)qtmp.data(), 4, n_per_row, imatrix, user_data); + repack_iq2_bn(4, n_per_row, qtmp.data(), qcur, false); + qcur += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq2_bn_r4(const block_iq2_bn * x, float * y, int64_t k) { + static_assert(QK_IQ1BN == 64); + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + const float * d4 = (const float *)x; + const uint8_t * qx = (const uint8_t *)(d4 + 4); + int nblock = n_per_row/QK_IQ1BN; + for (int ib = 0; ib < nblock; ++ib) { + for (int k = 0; k < 4; ++k) { + for (int l = 0; l < 4; ++l) for (int i = 0; i < 4; ++i) { + uint8_t q = qx[4*k + i + 16*l]; + y4[k][64*ib + 16*l + i + 0] = d4[k] * (((q >> 0) & 3) - 1); + y4[k][64*ib + 16*l + i + 4] = d4[k] * (((q >> 2) & 3) - 1); + y4[k][64*ib + 16*l + i + 8] = d4[k] * (((q >> 4) & 3) - 1); + y4[k][64*ib + 16*l + i + 12] = d4[k] * (((q >> 6) & 3) - 1); + } + } + qx += 64; + } +} + +void vec_dot_iq2_bn_r4_q8_K64(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_BN_R4, vx, 0, GGML_TYPE_Q8_K64, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q4_k_r4 +// + +void quantize_row_q4_k_r4_ref(const float * x, block_q4_k_r4 * y, int64_t k) { + quantize_q4_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_q4_k_r4(const float * x, void * y, int64_t k) { + quantize_q4_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t& d, uint8_t& m) { + if (j < 4) { + d = q[j] & 63; m = q[j + 4] & 63; + } else { + d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4); + m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4); + } +} +inline void convert_q4_k(const block_q4_K& x, uint8_t * L, uint8_t * Ld, uint8_t * Lm) { + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + get_scale_min_k4(2*ib64+0, x.scales, Ld[2*ib64+0], Lm[2*ib64+0]); + get_scale_min_k4(2*ib64+1, x.scales, Ld[2*ib64+1], Lm[2*ib64+1]); + for (int j = 0; j < 32; ++j) { + L[64*ib64+j+ 0] = x.qs[32*ib64+j] & 0xf; + L[64*ib64+j+32] = x.qs[32*ib64+j] >> 4; + } + } +} +} + +static void repack_q4_k(int nrows, int n_per_row, const block_q4_K * x, block_q4_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_q4_K * x4[4]; + uint8_t L[QK_K], Ld[QK_K/32], Lm[QK_K/32]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].scales_l, 0, QK_K/8); + std::memset(y[ibl].scales_h, 0, QK_K/16); + for (int k = 0; k < 4; ++k) { + y[ibl].d[k+0] = x4[k][ibl].d; + y[ibl].d[k+4] = x4[k][ibl].dmin; + convert_q4_k(x4[k][ibl], L, Ld, Lm); + for (int ib = 0; ib < QK_K/32; ++ib) { + y[ibl].scales_l[4*ib+k] = (Ld[ib] & 0xf) | ((Lm[ib] & 0xf) << 4); + uint8_t h = (Ld[ib] >> 4) | ((Lm[ib] >> 4) << 2); + y[ibl].scales_h[(4*ib+k)%16] |= (h << 4*((4*ib+k)/16)); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[64*ib+4*k+i+ 0] = L[32*ib+i+ 0] | (L[32*ib+i+ 8] << 4); + y[ibl].qs[64*ib+4*k+i+16] = L[32*ib+i+16] | (L[32*ib+i+24] << 4); + y[ibl].qs[64*ib+4*k+i+32] = L[32*ib+i+ 4] | (L[32*ib+i+12] << 4); + y[ibl].qs[64*ib+4*k+i+48] = L[32*ib+i+20] | (L[32*ib+i+28] << 4); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_q4_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_q4_K(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<32, block_q4_K, block_q4_k_r4, 4>(GGML_TYPE_Q4_K, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_q4_k); +} + +void dequantize_row_q4_k_r4(const block_q4_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k+0]); + const float m = GGML_FP16_TO_FP32(x[ibl].d[k+4]); + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 4*ib + k; + float dl = d * ((x[ibl].scales_l[is] & 0xf) | (((x[ibl].scales_h[is%16] >> 4*(is/16)) & 0x03) << 4)); + float ml = m * ((x[ibl].scales_l[is] >> 4) | (((x[ibl].scales_h[is%16] >> 4*(is/16)) & 0x0c) << 2)); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl * (x[ibl].qs[64*ib+4*k+i+ 0] & 0xf) - ml; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl * (x[ibl].qs[64*ib+4*k+i+ 0] >> 4) - ml; + y4[k][QK_K*ibl+32*ib+i+16] = dl * (x[ibl].qs[64*ib+4*k+i+16] & 0xf) - ml; + y4[k][QK_K*ibl+32*ib+i+24] = dl * (x[ibl].qs[64*ib+4*k+i+16] >> 4) - ml; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl * (x[ibl].qs[64*ib+4*k+i+32] & 0xf) - ml; + y4[k][QK_K*ibl+32*ib+i+12] = dl * (x[ibl].qs[64*ib+4*k+i+32] >> 4) - ml; + y4[k][QK_K*ibl+32*ib+i+20] = dl * (x[ibl].qs[64*ib+4*k+i+48] & 0xf) - ml; + y4[k][QK_K*ibl+32*ib+i+28] = dl * (x[ibl].qs[64*ib+4*k+i+48] >> 4) - ml; + } + } + } + } +} + +void vec_dot_q4_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q4_K_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q6_k_r4 +// + +void quantize_row_q6_k_r4_ref(const float * x, block_q6_k_r4 * y, int64_t k) { + quantize_q6_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_q6_k_r4(const float * x, void * y, int64_t k) { + quantize_q6_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +inline void convert_q6_k(const block_q6_K& x, uint8_t * L) { + const uint8_t * ql = x.ql; + const uint8_t * qh = x.qh; + + for (int n = 0; n < QK_K; n += 128) { + for (int l = 0; l < 32; ++l) { + L[n + l + 0] = (ql[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4); + L[n + l + 32] = (ql[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4); + L[n + l + 64] = (ql[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4); + L[n + l + 96] = (ql[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4); + } + ql += 64; + qh += 32; + } +} +} + +static void repack_q6_k(int nrows, int n_per_row, const block_q6_K * x, block_q6_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_q6_K * x4[4]; + uint8_t L[QK_K]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + convert_q6_k(x4[k][ibl], L); + for (int ib = 0; ib < QK_K/32; ++ib) { + y[ibl].scales[8*ib+k+0] = x4[k][ibl].scales[2*ib+0]; + y[ibl].scales[8*ib+k+4] = x4[k][ibl].scales[2*ib+1]; + for (int i = 0; i < 4; ++i) { + y[ibl].ql[64*ib+4*k+i+ 0] = (L[32*ib+i+ 0] & 0xf) | ((L[32*ib+i+ 8] & 0xf) << 4); + y[ibl].ql[64*ib+4*k+i+16] = (L[32*ib+i+16] & 0xf) | ((L[32*ib+i+24] & 0xf) << 4); + y[ibl].ql[64*ib+4*k+i+32] = (L[32*ib+i+ 4] & 0xf) | ((L[32*ib+i+12] & 0xf) << 4); + y[ibl].ql[64*ib+4*k+i+48] = (L[32*ib+i+20] & 0xf) | ((L[32*ib+i+28] & 0xf) << 4); + y[ibl].qh[32*ib+4*k+i+ 0] = (L[32*ib+i+ 0] >> 4) | ((L[32*ib+i+ 8] >> 4) << 2) | ((L[32*ib+i+ 4] >> 4) << 4) | ((L[32*ib+i+12] >> 4) << 6); + y[ibl].qh[32*ib+4*k+i+16] = (L[32*ib+i+16] >> 4) | ((L[32*ib+i+24] >> 4) << 2) | ((L[32*ib+i+20] >> 4) << 4) | ((L[32*ib+i+28] >> 4) << 6); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_q6_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_q6_K(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<16, block_q6_K, block_q6_k_r4, 4>(GGML_TYPE_Q6_K, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_q6_k); +} + +void dequantize_row_q6_k_r4(const block_q6_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + auto ql = x[ibl].ql; + auto qh = x[ibl].qh; + for (int ib = 0; ib < QK_K/32; ++ib) { + float dl1 = d * x[ibl].scales[8*ib+k+0]; + float dl2 = d * x[ibl].scales[8*ib+k+4]; + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl1 * (((ql[4*k+i+ 0] & 0xf) | ((qh[4*k+i+ 0] << 4) & 0x30)) - 32); + y4[k][QK_K*ibl+32*ib+i+ 8] = dl1 * (((ql[4*k+i+ 0] >> 4) | ((qh[4*k+i+ 0] << 2) & 0x30)) - 32); + y4[k][QK_K*ibl+32*ib+i+16] = dl2 * (((ql[4*k+i+16] & 0xf) | ((qh[4*k+i+16] << 4) & 0x30)) - 32); + y4[k][QK_K*ibl+32*ib+i+24] = dl2 * (((ql[4*k+i+16] >> 4) | ((qh[4*k+i+16] << 2) & 0x30)) - 32); + y4[k][QK_K*ibl+32*ib+i+ 4] = dl1 * (((ql[4*k+i+32] & 0xf) | ((qh[4*k+i+ 0] >> 0) & 0x30)) - 32); + y4[k][QK_K*ibl+32*ib+i+12] = dl1 * (((ql[4*k+i+32] >> 4) | ((qh[4*k+i+ 0] >> 2) & 0x30)) - 32); + y4[k][QK_K*ibl+32*ib+i+20] = dl2 * (((ql[4*k+i+48] & 0xf) | ((qh[4*k+i+16] >> 0) & 0x30)) - 32); + y4[k][QK_K*ibl+32*ib+i+28] = dl2 * (((ql[4*k+i+48] >> 4) | ((qh[4*k+i+16] >> 2) & 0x30)) - 32); + } + ql += 64; + qh += 32; + } + } + } +} + +void vec_dot_q6_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q6_K_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + + +// +// ========================================= q5_k_r4 +// + +void quantize_row_q5_k_r4_ref(const float * x, block_q5_k_r4 * y, int64_t k) { + quantize_q5_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_q5_k_r4(const float * x, void * y, int64_t k) { + quantize_q5_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +inline void convert_q5_k(const block_q5_K& x, uint8_t * L, uint8_t * Ld, uint8_t * Lm) { + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + get_scale_min_k4(2*ib64+0, x.scales, Ld[2*ib64+0], Lm[2*ib64+0]); + get_scale_min_k4(2*ib64+1, x.scales, Ld[2*ib64+1], Lm[2*ib64+1]); + for (int j = 0; j < 32; ++j) { + L[64*ib64+j+ 0] = (x.qs[32*ib64+j] & 0xf) | (((x.qh[j] >> (2*ib64+0)) & 1) << 4); + L[64*ib64+j+32] = (x.qs[32*ib64+j] >> 4) | (((x.qh[j] >> (2*ib64+1)) & 1) << 4); + } + } +} +} + +static void repack_q5_k(int nrows, int n_per_row, const block_q5_K * x, block_q5_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_q5_K * x4[4]; + uint8_t L[QK_K], Ld[QK_K/32], Lm[QK_K/32]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].scales_l, 0, QK_K/8); + std::memset(y[ibl].scales_h, 0, QK_K/16); + for (int k = 0; k < 4; ++k) { + y[ibl].d[k+0] = x4[k][ibl].d; + y[ibl].d[k+4] = x4[k][ibl].dmin; + convert_q5_k(x4[k][ibl], L, Ld, Lm); + for (int ib = 0; ib < QK_K/32; ++ib) { + y[ibl].scales_l[4*ib+k] = (Ld[ib] & 0xf) | ((Lm[ib] & 0xf) << 4); + uint8_t h = (Ld[ib] >> 4) | ((Lm[ib] >> 4) << 2); + y[ibl].scales_h[(4*ib+k)%16] |= (h << 4*((4*ib+k)/16)); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[64*ib+4*k+i+ 0] = (L[32*ib+i+ 0] & 0xf) | ((L[32*ib+i+ 8] & 0xf) << 4); + y[ibl].qs[64*ib+4*k+i+16] = (L[32*ib+i+16] & 0xf) | ((L[32*ib+i+24] & 0xf) << 4); + y[ibl].qs[64*ib+4*k+i+32] = (L[32*ib+i+ 4] & 0xf) | ((L[32*ib+i+12] & 0xf) << 4); + y[ibl].qs[64*ib+4*k+i+48] = (L[32*ib+i+20] & 0xf) | ((L[32*ib+i+28] & 0xf) << 4); + y[ibl].qh[16*ib+4*k+i+ 0] = ((L[32*ib+i+ 0] >> 4) << 0) | ((L[32*ib+i+ 8] >> 4) << 1) | ((L[32*ib+i+ 4] >> 4) << 2) | ((L[32*ib+i+12] >> 4) << 3) | + ((L[32*ib+i+16] >> 4) << 4) | ((L[32*ib+i+24] >> 4) << 5) | ((L[32*ib+i+20] >> 4) << 6) | ((L[32*ib+i+28] >> 4) << 7); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_q5_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_q5_K(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<32, block_q5_K, block_q5_k_r4, 4>(GGML_TYPE_Q5_K, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_q5_k); +} + +void dequantize_row_q5_k_r4(const block_q5_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k+0]); + const float m = GGML_FP16_TO_FP32(x[ibl].d[k+4]); + auto ql = x[ibl].qs; + auto qh = x[ibl].qh; + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 4*ib + k; + float dl = d * ((x[ibl].scales_l[is] & 0xf) | (((x[ibl].scales_h[is%16] >> 4*(is/16)) & 0x03) << 4)); + float ml = m * ((x[ibl].scales_l[is] >> 4) | (((x[ibl].scales_h[is%16] >> 4*(is/16)) & 0x0c) << 2)); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl * ((ql[4*k+i+ 0] & 0xf) | ((qh[4*k+i] << 4) & 0x10)) - ml; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl * ((ql[4*k+i+ 0] >> 4) | ((qh[4*k+i] << 3) & 0x10)) - ml; + y4[k][QK_K*ibl+32*ib+i+16] = dl * ((ql[4*k+i+16] & 0xf) | ((qh[4*k+i] >> 0) & 0x10)) - ml; + y4[k][QK_K*ibl+32*ib+i+24] = dl * ((ql[4*k+i+16] >> 4) | ((qh[4*k+i] >> 1) & 0x10)) - ml; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl * ((ql[4*k+i+32] & 0xf) | ((qh[4*k+i] << 2) & 0x10)) - ml; + y4[k][QK_K*ibl+32*ib+i+12] = dl * ((ql[4*k+i+32] >> 4) | ((qh[4*k+i] << 1) & 0x10)) - ml; + y4[k][QK_K*ibl+32*ib+i+20] = dl * ((ql[4*k+i+48] & 0xf) | ((qh[4*k+i] >> 2) & 0x10)) - ml; + y4[k][QK_K*ibl+32*ib+i+28] = dl * ((ql[4*k+i+48] >> 4) | ((qh[4*k+i] >> 3) & 0x10)) - ml; + } + ql += 64; + qh += 16; + } + } + } +} + +void vec_dot_q5_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q5_K_R4, vx, 0, GGML_TYPE_Q8_K32, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q3_k_r4 +// + +void quantize_row_q3_k_r4_ref(const float * x, block_q3_k_r4 * y, int64_t k) { + quantize_q3_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_q3_k_r4(const float * x, void * y, int64_t k) { + quantize_q3_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +inline void convert_q3_k(const block_q3_K& x, uint8_t * L, uint8_t * Ld) { + constexpr uint32_t kmask1 = 0x03030303; + constexpr uint32_t kmask2 = 0x0f0f0f0f; + uint32_t aux[4]; + memcpy(aux, x.scales, 12); + uint32_t tmp = aux[2]; + aux[2] = ((aux[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4); + aux[3] = ((aux[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4); + aux[0] = (aux[0] & kmask2) | (((tmp >> 0) & kmask1) << 4); + aux[1] = (aux[1] & kmask2) | (((tmp >> 2) & kmask1) << 4); + std::memcpy(Ld, aux, 16); + + const uint8_t * q = x.qs; + const uint8_t * hm = x.hmask; + uint8_t m = 1; + for (int n = 0; n < QK_K; n += 128) { + int shift = 0; + for (int j = 0; j < 4; ++j) { + for (int l = 0; l < 32; ++l) { + *L++ = ((q[l] >> shift) & 3) + ((hm[l] & m) ? 4 : 0); + } + shift += 2; + m <<= 1; + } + q += 32; + } +} +} + +static void repack_q3_k(int nrows, int n_per_row, const block_q3_K * x, block_q3_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_q3_K * x4[4]; + uint8_t L[QK_K], Ld[QK_K/16]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].scales_l, 0, QK_K/8); + std::memset(y[ibl].scales_h, 0, QK_K/16); + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + convert_q3_k(x4[k][ibl], L, Ld); + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 8*ib+k; + y[ibl].scales_l[is%32] |= (Ld[2*ib+0] & 0xf) << 4*(is/32); + y[ibl].scales_h[is%16] |= (Ld[2*ib+0] >> 4) << 2*(is/16); + is += 4; + y[ibl].scales_l[is%32] |= (Ld[2*ib+1] & 0xf) << 4*(is/32); + y[ibl].scales_h[is%16] |= (Ld[2*ib+1] >> 4) << 2*(is/16); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[32*ib+4*k+i+ 0] = ((L[32*ib+i+ 0] & 0x3) << 0) | ((L[32*ib+i+ 4] & 0x3) << 2) | ((L[32*ib+i+ 8] & 0x3) << 4) | ((L[32*ib+i+12] & 0x3) << 6); + y[ibl].qs[32*ib+4*k+i+16] = ((L[32*ib+i+16] & 0x3) << 0) | ((L[32*ib+i+20] & 0x3) << 2) | ((L[32*ib+i+24] & 0x3) << 4) | ((L[32*ib+i+28] & 0x3) << 6); + y[ibl].qh[16*ib+4*k+i+ 0] = ((L[32*ib+i+ 0] >> 2) << 0) | ((L[32*ib+i+ 4] >> 2) << 1) | ((L[32*ib+i+ 8] >> 2) << 2) | ((L[32*ib+i+12] >> 2) << 3) + | ((L[32*ib+i+16] >> 2) << 4) | ((L[32*ib+i+20] >> 2) << 5) | ((L[32*ib+i+24] >> 2) << 6) | ((L[32*ib+i+28] >> 2) << 7); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_q3_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_q3_K(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<16, block_q3_K, block_q3_k_r4, 4>(GGML_TYPE_Q3_K, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_q3_k); +} + +void dequantize_row_q3_k_r4(const block_q3_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + auto ql = x[ibl].qs; + auto qh = x[ibl].qh; + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 8*ib + k; + float dl1 = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 0x03) << 4)) - 32); + is += 4; + float dl2 = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 0x03) << 4)) - 32); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl1 * ((((ql[4*k+i+ 0] >> 0) & 3) | ((qh[4*k+i] << 2) & 4)) - 4); + y4[k][QK_K*ibl+32*ib+i+ 4] = dl1 * ((((ql[4*k+i+ 0] >> 2) & 3) | ((qh[4*k+i] << 1) & 4)) - 4); + y4[k][QK_K*ibl+32*ib+i+ 8] = dl1 * ((((ql[4*k+i+ 0] >> 4) & 3) | ((qh[4*k+i] << 0) & 4)) - 4); + y4[k][QK_K*ibl+32*ib+i+12] = dl1 * ((((ql[4*k+i+ 0] >> 6) & 3) | ((qh[4*k+i] >> 1) & 4)) - 4); + y4[k][QK_K*ibl+32*ib+i+16] = dl2 * ((((ql[4*k+i+16] >> 0) & 3) | ((qh[4*k+i] >> 2) & 4)) - 4); + y4[k][QK_K*ibl+32*ib+i+20] = dl2 * ((((ql[4*k+i+16] >> 2) & 3) | ((qh[4*k+i] >> 3) & 4)) - 4); + y4[k][QK_K*ibl+32*ib+i+24] = dl2 * ((((ql[4*k+i+16] >> 4) & 3) | ((qh[4*k+i] >> 4) & 4)) - 4); + y4[k][QK_K*ibl+32*ib+i+28] = dl2 * ((((ql[4*k+i+16] >> 6) & 3) | ((qh[4*k+i] >> 5) & 4)) - 4); + } + ql += 32; + qh += 16; + } + } + } +} + +void vec_dot_q3_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q3_K_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q2_k_r4 +// + +void quantize_row_q2_k_r4_ref(const float * x, block_q2_k_r4 * y, int64_t k) { + quantize_q3_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_q2_k_r4(const float * x, void * y, int64_t k) { + quantize_q2_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +inline void convert_q2_k(const block_q2_K& x, uint8_t * L) { + + const uint8_t * qs = x.qs; + for (int n = 0; n < QK_K; n += 128) { + for (int j = 0; j < 32; ++j) { + L[n + j + 0] = (qs[j] >> 0) & 0x3; + L[n + j + 32] = (qs[j] >> 2) & 0x3; + L[n + j + 64] = (qs[j] >> 4) & 0x3; + L[n + j + 96] = (qs[j] >> 6) & 0x3; + } + qs += 32; + } +} +} + +static void repack_q2_k(int nrows, int n_per_row, const block_q2_K * x, block_q2_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_q2_K * x4[4]; + uint8_t L[QK_K]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + y[ibl].d[k+0] = x4[k][ibl].d; + y[ibl].d[k+4] = x4[k][ibl].dmin; + for (int ib = 0; ib < QK_K/16; ++ib) { + y[ibl].scales[4*ib+k] = x4[k][ibl].scales[ib]; + } + convert_q2_k(x4[k][ibl], L); + for (int ib = 0; ib < QK_K/32; ++ib) { + for (int i = 0; i < 4; ++i) { + y[ibl].qs[32*ib+4*k+i+ 0] = ((L[32*ib+i+ 0] & 0x3) << 0) | ((L[32*ib+i+ 4] & 0x3) << 2) | ((L[32*ib+i+ 8] & 0x3) << 4) | ((L[32*ib+i+12] & 0x3) << 6); + y[ibl].qs[32*ib+4*k+i+16] = ((L[32*ib+i+16] & 0x3) << 0) | ((L[32*ib+i+20] & 0x3) << 2) | ((L[32*ib+i+24] & 0x3) << 4) | ((L[32*ib+i+28] & 0x3) << 6); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_q2_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_q2_K(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<16, block_q2_K, block_q2_k_r4, 4>(GGML_TYPE_Q2_K, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_q2_k); +} + +void dequantize_row_q2_k_r4(const block_q2_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k+0]); + const float m = GGML_FP16_TO_FP32(x[ibl].d[k+4]); + auto ql = x[ibl].qs; + for (int ib = 0; ib < QK_K/32; ++ib) { + float dl1 = d * (x[ibl].scales[8*ib + k + 0] & 0xf); + float ml1 = m * (x[ibl].scales[8*ib + k + 0] >> 4); + float dl2 = d * (x[ibl].scales[8*ib + k + 4] & 0xf); + float ml2 = m * (x[ibl].scales[8*ib + k + 4] >> 4); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl1 * ((ql[4*k+i+ 0] >> 0) & 3) - ml1; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl1 * ((ql[4*k+i+ 0] >> 2) & 3) - ml1; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl1 * ((ql[4*k+i+ 0] >> 4) & 3) - ml1; + y4[k][QK_K*ibl+32*ib+i+12] = dl1 * ((ql[4*k+i+ 0] >> 6) & 3) - ml1; + y4[k][QK_K*ibl+32*ib+i+16] = dl2 * ((ql[4*k+i+16] >> 0) & 3) - ml2; + y4[k][QK_K*ibl+32*ib+i+20] = dl2 * ((ql[4*k+i+16] >> 2) & 3) - ml2; + y4[k][QK_K*ibl+32*ib+i+24] = dl2 * ((ql[4*k+i+16] >> 4) & 3) - ml2; + y4[k][QK_K*ibl+32*ib+i+28] = dl2 * ((ql[4*k+i+16] >> 6) & 3) - ml2; + } + ql += 32; + } + } + } +} + +void vec_dot_q2_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q2_K_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq4_k_r4 +// + +void quantize_row_iq4_k_r4_ref(const float * x, block_iq4_k_r4 * y, int64_t k) { + quantize_iq4_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq4_k_r4(const float * x, void * y, int64_t k) { + quantize_iq4_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static void repack_iq4_k(int nrows, int n_per_row, const block_iq4_k * x, block_iq4_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq4_k * x4[4]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].extra, 0, 8); + std::memset(y[ibl].scales_l, 0, QK_K/8); + std::memset(y[ibl].scales_h, 0, QK_K/16); + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + auto extra = x4[k][ibl].extra; + for (int ib = 0; ib < QK_K/32; ++ib) { + if (extra & 1) y[ibl].extra[k+0] |= (1 << ib); + if (extra & 2) y[ibl].extra[k+4] |= (1 << ib); + extra >>= 2; + uint8_t sl1 = x4[k][ibl].scales_l[ib] & 0xf; + uint8_t sl2 = x4[k][ibl].scales_l[ib] >> 4; + uint8_t sh = x4[k][ibl].scales_h[ib/2] >> 4*(ib%2); + uint8_t sh1 = (sh >> 0) & 3; + uint8_t sh2 = (sh >> 2) & 3; + int i = 8*ib + k; + y[ibl].scales_l[i%32] |= (sl1 << 4*(i/32)); + y[ibl].scales_h[i%16] |= (sh1 << 2*(i/16)); + i += 4; + y[ibl].scales_l[i%32] |= (sl2 << 4*(i/32)); + y[ibl].scales_h[i%16] |= (sh2 << 2*(i/16)); + } + } + for (int ib = 0; ib < QK_K/32; ++ib) { + for (int k = 0; k < 4; ++k) for (int i = 0; i < 4; ++i) { + y[ibl].qs[64*ib+4*k+i+ 0] = (x4[k][ibl].qs[16*ib+i+0] & 0xf) | ((x4[k][ibl].qs[16*ib+i+ 8] & 0x0f) << 4); // 0....3 + 8...11 from each row + y[ibl].qs[64*ib+4*k+i+16] = (x4[k][ibl].qs[16*ib+i+0] >> 4) | ((x4[k][ibl].qs[16*ib+i+ 8] & 0xf0)); // 16...19 + 24...27 from each row + y[ibl].qs[64*ib+4*k+i+32] = (x4[k][ibl].qs[16*ib+i+4] & 0xf) | ((x4[k][ibl].qs[16*ib+i+12] & 0x0f) << 4); // 4....7 + 12...15 from each row + y[ibl].qs[64*ib+4*k+i+48] = (x4[k][ibl].qs[16*ib+i+4] >> 4) | ((x4[k][ibl].qs[16*ib+i+12] & 0xf0)); // 20...23 + 28...31 from each row + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq4_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + char * qcur = (char *)dst; + auto row_size = ggml_row_size(GGML_TYPE_IQ4_K, n_per_row); + std::vector qtmp(4*row_size); + for (int row = 0; row < nrows; row += 4) { + quantize_iq4_k(src, (void *)qtmp.data(), 4, n_per_row, imatrix, user_data); + repack_iq4_k(4, n_per_row, (const block_iq4_k *)qtmp.data(), (block_iq4_k_r4 *)qcur, false); + qcur += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq4_k_r4(const block_iq4_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 8*ib + k; + float dl1 = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 3) << 4)) - 32); + is += 4; + float dl2 = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 3) << 4)) - 32); + auto values1 = iq4k_values + (x[ibl].extra[k+0] & (1 << ib) ? 16 : 0); + auto values2 = iq4k_values + (x[ibl].extra[k+4] & (1 << ib) ? 16 : 0); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl1 * values1[x[ibl].qs[64*ib+4*k+i+ 0] & 0xf]; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl1 * values1[x[ibl].qs[64*ib+4*k+i+ 0] >> 4]; + y4[k][QK_K*ibl+32*ib+i+16] = dl2 * values2[x[ibl].qs[64*ib+4*k+i+16] & 0xf]; + y4[k][QK_K*ibl+32*ib+i+24] = dl2 * values2[x[ibl].qs[64*ib+4*k+i+16] >> 4]; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl1 * values1[x[ibl].qs[64*ib+4*k+i+32] & 0xf]; + y4[k][QK_K*ibl+32*ib+i+12] = dl1 * values1[x[ibl].qs[64*ib+4*k+i+32] >> 4]; + y4[k][QK_K*ibl+32*ib+i+20] = dl2 * values2[x[ibl].qs[64*ib+4*k+i+48] & 0xf]; + y4[k][QK_K*ibl+32*ib+i+28] = dl2 * values2[x[ibl].qs[64*ib+4*k+i+48] >> 4]; + } + } + } + } +} + +void vec_dot_iq4_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ4_K_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq5_k_r4 +// + +void quantize_row_iq5_k_r4_ref(const float * x, block_iq5_k_r4 * y, int64_t k) { + quantize_iq5_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq5_k_r4(const float * x, void * y, int64_t k) { + quantize_iq5_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +template +inline void convert_iq5_k(const Block& x, uint8_t * L) { + const uint8_t * qs = x.qs; + const uint8_t * qh = x.qh; + int shift = 0; + for (int ib64 = 0; ib64 < QK_K/64; ++ib64) { + for (int j = 0; j < 16; ++j) { + L[j+ 0] = (qs[j+ 0] & 0xf) | (((qh[j+ 0] >> shift) & 1) << 4); + L[j+16] = (qs[j+16] & 0xf) | (((qh[j+16] >> shift) & 1) << 4); + L[j+32] = (qs[j+ 0] >> 4) | (((qh[j+ 0] >> shift) & 2) << 3); + L[j+48] = (qs[j+16] >> 4) | (((qh[j+16] >> shift) & 2) << 3); + } + L += 64; + qs += 32; + shift += 2; + if (shift == 8) { qh += 32; shift = 0; } + } +} +} + +static void repack_iq5_k(int nrows, int n_per_row, const block_iq5_k * x, block_iq5_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq5_k * x4[4]; + uint8_t L[QK_K]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].extra, 0, 8); + std::memset(y[ibl].scales_l, 0, QK_K/8); + std::memset(y[ibl].scales_h, 0, QK_K/16); + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + auto extra = x4[k][ibl].extra; + convert_iq5_k(x4[k][ibl], L); + for (int ib = 0; ib < QK_K/32; ++ib) { + if (extra & 1) y[ibl].extra[k+0] |= (1 << ib); + if (extra & 2) y[ibl].extra[k+4] |= (1 << ib); + extra >>= 2; + uint8_t sl1 = x4[k][ibl].scales_l[ib] & 0xf; + uint8_t sl2 = x4[k][ibl].scales_l[ib] >> 4; + uint8_t sh = x4[k][ibl].scales_h[ib/2] >> 4*(ib%2); + uint8_t sh1 = (sh >> 0) & 3; + uint8_t sh2 = (sh >> 2) & 3; + int i = 8*ib + k; + y[ibl].scales_l[i%32] |= (sl1 << 4*(i/32)); + y[ibl].scales_h[i%16] |= (sh1 << 2*(i/16)); + i += 4; + y[ibl].scales_l[i%32] |= (sl2 << 4*(i/32)); + y[ibl].scales_h[i%16] |= (sh2 << 2*(i/16)); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[64*ib+4*k+i+ 0] = (L[32*ib+i+ 0] & 0xf) | ((L[32*ib+i+ 8] & 0xf) << 4); // 0....3 + 8...11 from each row + y[ibl].qs[64*ib+4*k+i+16] = (L[32*ib+i+16] & 0xf) | ((L[32*ib+i+24] & 0xf) << 4); // 16...19 + 24...27 from each row + y[ibl].qs[64*ib+4*k+i+32] = (L[32*ib+i+ 4] & 0xf) | ((L[32*ib+i+12] & 0xf) << 4); // 4....7 + 12...15 from each row + y[ibl].qs[64*ib+4*k+i+48] = (L[32*ib+i+20] & 0xf) | ((L[32*ib+i+28] & 0xf) << 4); // 20...23 + 28...31 from each row + y[ibl].qh[16*ib+4*k+i ] = ((L[32*ib+i+ 0] >> 4) << 0) | ((L[32*ib+i+ 8] >> 4) << 1) | ((L[32*ib+i+16] >> 4) << 2) | ((L[32*ib+i+24] >> 4) << 3) + | ((L[32*ib+i+ 4] >> 4) << 4) | ((L[32*ib+i+12] >> 4) << 5) | ((L[32*ib+i+20] >> 4) << 6) | ((L[32*ib+i+28] >> 4) << 7); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq5_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + char * qcur = (char *)dst; + auto row_size = ggml_row_size(GGML_TYPE_IQ5_K, n_per_row); + std::vector qtmp(4*row_size); + for (int row = 0; row < nrows; row += 4) { + quantize_iq5_k(src, (void *)qtmp.data(), 4, n_per_row, imatrix, user_data); + repack_iq5_k(4, n_per_row, (const block_iq5_k *)qtmp.data(), (block_iq5_k_r4 *)qcur, false); + qcur += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq5_k_r4(const block_iq5_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 8*ib + k; + float dl1 = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 3) << 4)) - 32); + is += 4; + float dl2 = d * ((((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) | (((x[ibl].scales_h[is%16] >> 2*(is/16)) & 3) << 4)) - 32); + auto values1 = iq5nl_values + (x[ibl].extra[k+0] & (1 << ib) ? 32 : 0); + auto values2 = iq5nl_values + (x[ibl].extra[k+4] & (1 << ib) ? 32 : 0); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl1 * values1[(x[ibl].qs[64*ib+4*k+i+ 0] & 0xf) | (((x[ibl].qh[16*ib+4*k+i] >> 0) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl1 * values1[(x[ibl].qs[64*ib+4*k+i+ 0] >> 4) | (((x[ibl].qh[16*ib+4*k+i] >> 1) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+16] = dl2 * values2[(x[ibl].qs[64*ib+4*k+i+16] & 0xf) | (((x[ibl].qh[16*ib+4*k+i] >> 2) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+24] = dl2 * values2[(x[ibl].qs[64*ib+4*k+i+16] >> 4) | (((x[ibl].qh[16*ib+4*k+i] >> 3) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl1 * values1[(x[ibl].qs[64*ib+4*k+i+32] & 0xf) | (((x[ibl].qh[16*ib+4*k+i] >> 4) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+12] = dl1 * values1[(x[ibl].qs[64*ib+4*k+i+32] >> 4) | (((x[ibl].qh[16*ib+4*k+i] >> 5) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+20] = dl2 * values2[(x[ibl].qs[64*ib+4*k+i+48] & 0xf) | (((x[ibl].qh[16*ib+4*k+i] >> 6) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+28] = dl2 * values2[(x[ibl].qs[64*ib+4*k+i+48] >> 4) | (((x[ibl].qh[16*ib+4*k+i] >> 7) & 1) << 4)]; + } + } + } + } +} + +void vec_dot_iq5_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ5_K_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq5_ks_r4 +// + +void quantize_row_iq5_ks_r4_ref(const float * x, block_iq5_ks_r4 * y, int64_t k) { + quantize_iq5_ks_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq5_ks_r4(const float * x, void * y, int64_t k) { + quantize_iq5_ks_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static void repack_iq5_ks(int nrows, int n_per_row, const block_iq5_ks * x, block_iq5_ks_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ5_KS, n_per_row); + int nblock = n_per_row/QK_K; + const block_iq5_ks * x4[4]; + uint8_t L[QK_K]; + char * cy = (char *)y; + const char * cx = (const char *)x; + for (int row = 0; row < nrows; row += 4) { + float * dptr = (float *)cy; + block_iq5_ks_r4 * y = (block_iq5_ks_r4 *)(dptr + 4); + for (int k = 0; k < 4; ++k) { + auto dk = (const float *)(cx + k*row_size); + dptr[k] = dk[0]; + x4[k] = (const block_iq5_ks *)(dk + 1); + } + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + convert_iq5_k(x4[k][ibl], L); + for (int ib = 0; ib < QK_K/32; ++ib) { + y[ibl].scales[4*ib+k] = x4[k][ibl].scales[ib]; + for (int i = 0; i < 4; ++i) { + y[ibl].qs[64*ib+4*k+i+ 0] = (L[32*ib+i+ 0] & 0xf) | ((L[32*ib+i+ 8] & 0xf) << 4); // 0....3 + 8...11 from each row + y[ibl].qs[64*ib+4*k+i+16] = (L[32*ib+i+16] & 0xf) | ((L[32*ib+i+24] & 0xf) << 4); // 16...19 + 24...27 from each row + y[ibl].qs[64*ib+4*k+i+32] = (L[32*ib+i+ 4] & 0xf) | ((L[32*ib+i+12] & 0xf) << 4); // 4....7 + 12...15 from each row + y[ibl].qs[64*ib+4*k+i+48] = (L[32*ib+i+20] & 0xf) | ((L[32*ib+i+28] & 0xf) << 4); // 20...23 + 28...31 from each row + y[ibl].qh[16*ib+4*k+i ] = ((L[32*ib+i+ 0] >> 4) << 0) | ((L[32*ib+i+ 8] >> 4) << 1) | ((L[32*ib+i+16] >> 4) << 2) | ((L[32*ib+i+24] >> 4) << 3) + | ((L[32*ib+i+ 4] >> 4) << 4) | ((L[32*ib+i+12] >> 4) << 5) | ((L[32*ib+i+20] >> 4) << 6) | ((L[32*ib+i+28] >> 4) << 7); + } + } + } + } + cx += 4*row_size; + cy += 4*row_size; + } +} + +size_t quantize_iq5_ks_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + char * qcur = (char *)dst; + auto row_size = ggml_row_size(GGML_TYPE_IQ5_KS, n_per_row); + std::vector qtmp(4*row_size); + for (int row = 0; row < nrows; row += 4) { + quantize_iq5_ks(src, (void *)qtmp.data(), 4, n_per_row, imatrix, user_data); + repack_iq5_ks(4, n_per_row, (const block_iq5_ks *)qtmp.data(), (block_iq5_ks_r4 *)qcur, false); + qcur += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq5_ks_r4(const block_iq5_ks_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + //auto row_size = ggml_row_size(GGML_TYPE_IQ5_KS, n_per_row); + int nblock = n_per_row/QK_K; + const float * dptr = (const float *)x; + x = (const block_iq5_ks_r4 *)(dptr + 4); + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = dptr[k]; + //if (!isfinite(d)) { + // printf("Oops: d = %g for ibl = %d, k = %d\n", d, ibl, k); exit(1); + //} + for (int ib = 0; ib < QK_K/32; ++ib) { + uint8_t sc = x[ibl].scales[4*ib+k]; + float dl = d * ((sc & 254) - 127); + //if (!isfinite(dl)) { + // printf("Oops: dl = %g for ibl = %d, k = %d, ib = %d, d = %g, sc = %u\n", dl, ibl, k, ib, d, sc); exit(1); + //} + auto values = iq5nl_values + ((sc & 1) << 5); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl * values[(x[ibl].qs[64*ib+4*k+i+ 0] & 0xf) | (((x[ibl].qh[16*ib+4*k+i] >> 0) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl * values[(x[ibl].qs[64*ib+4*k+i+ 0] >> 4) | (((x[ibl].qh[16*ib+4*k+i] >> 1) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+16] = dl * values[(x[ibl].qs[64*ib+4*k+i+16] & 0xf) | (((x[ibl].qh[16*ib+4*k+i] >> 2) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+24] = dl * values[(x[ibl].qs[64*ib+4*k+i+16] >> 4) | (((x[ibl].qh[16*ib+4*k+i] >> 3) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl * values[(x[ibl].qs[64*ib+4*k+i+32] & 0xf) | (((x[ibl].qh[16*ib+4*k+i] >> 4) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+12] = dl * values[(x[ibl].qs[64*ib+4*k+i+32] >> 4) | (((x[ibl].qh[16*ib+4*k+i] >> 5) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+20] = dl * values[(x[ibl].qs[64*ib+4*k+i+48] & 0xf) | (((x[ibl].qh[16*ib+4*k+i] >> 6) & 1) << 4)]; + y4[k][QK_K*ibl+32*ib+i+28] = dl * values[(x[ibl].qs[64*ib+4*k+i+48] >> 4) | (((x[ibl].qh[16*ib+4*k+i] >> 7) & 1) << 4)]; + } + //for (int i = 0; i < 32; ++i) { + // if (!isfinite(y4[k][QK_K*ibl+32*ib+i])) { + // printf("Oops: y4[%d][%d, %d, %d] = %g\n", k, ibl, ib, i, y4[k][QK_K*ibl+32*ib+i]); + // printf("d = %g, dl = %g\n", d, dl); + // exit(1); + // } + //} + } + } + } +} + +void vec_dot_iq5_ks_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ5_KS_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q8_k_r8 +// + +void quantize_row_q8_k_r8_ref(const float * x, block_q8_k_r8 * y, int64_t k) { + quantize_q8_k_r8(x, (void *)y, 8, k/8, nullptr, nullptr); +} + +void quantize_row_q8_k_r8(const float * x, void * y, int64_t k) { + quantize_q8_k_r8(x, y, 8, k/8, nullptr, nullptr); +} + +static void repack_q8_k(int nrows, int n_per_row, const block_q8_K * x, block_q8_k_r8 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%8 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_q8_K * x8[8]; + for (int row = 0; row < nrows; row += 8) { + for (int k = 0; k < 8; ++k) x8[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 8; ++k) { + y[ibl].d[k] = GGML_FP32_TO_FP16(x8[k][ibl].d); + for (int ib = 0; ib < QK_K/4; ++ib) { + for (int i = 0; i < 4; ++i) y[ibl].qs[32*ib + 4*k + i] = x8[k][ibl].qs[4*ib+i]; + } + } +#ifdef HAVE_FANCY_SIMD + if (online) { + for (int l = 0; l < 32; ++l) { + auto v = _mm512_xor_si512(_mm512_loadu_si512((const __m512i *)y[ibl].qs + l), _mm512_set1_epi8(-128)); + _mm512_storeu_si512((__m512i *)y[ibl].qs + l, v); + } + } +#endif + } + x += 8*nblock; + y += nblock; + } +} +#ifdef HAVE_FANCY_SIMD +static void modify_q8_k_r8(int64_t k, char * cy) { + auto y = (block_q8_k_r8 *)cy; + int nb = k/(256*8); + for (int ib = 0; ib < nb; ++ib) { + for (int l = 0; l < 32; ++l) { + auto v = _mm512_xor_si512(_mm512_loadu_si512((const __m512i *)y[ib].qs + l), _mm512_set1_epi8(-128)); + _mm512_storeu_si512((__m512i *)y[ib].qs + l, v); + } + } +} +#endif + +size_t quantize_q8_k_r8(const float * src, void * dst, int64_t nrows, int64_t n_per_row, [[maybe_unused]] const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%8 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + char * qcur = (char *)dst; + auto row_size_0 = ggml_row_size(GGML_TYPE_Q8_K, n_per_row); + auto row_size_1 = ggml_row_size(GGML_TYPE_Q8_K_R8, n_per_row); + std::vector qtmp(8*row_size_0); + for (int row = 0; row < nrows; row += 8) { + quantize_row_q8_K32(src, (void *)qtmp.data(), 8*n_per_row); + repack_q8_k(8, n_per_row, (const block_q8_K *)qtmp.data(), (block_q8_k_r8 *)qcur, false); + qcur += 8*row_size_1; + src += 8*n_per_row; + } + return nrows*row_size_1; +} + +void dequantize_row_q8_k_r8(const block_q8_k_r8 * x, float * y, int64_t k) { + auto n_per_row = k/8; + float * y8[8]; + for (int k = 0; k < 8; ++k) y8[k] = y + n_per_row*k; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 8; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/4; ++ib) { + for (int i = 0; i < 4; ++i) { + y8[k][QK_K*ibl+4*ib+i] = d * x[ibl].qs[32*ib+4*k+i]; + } + } + } + } +} + +void vec_dot_q8_k_r8_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q8_K_R8, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q8_k_r16 +// + +void quantize_row_q8_k_r16_ref(const float * x, block_q8_k_r16 * y, int64_t k) { + quantize_q8_k_r16(x, (void *)y, 16, k/16, nullptr, nullptr); +} + +void quantize_row_q8_k_r16(const float * x, void * y, int64_t k) { + quantize_q8_k_r16(x, y, 16, k/16, nullptr, nullptr); +} + +static void repack_q16_k(int nrows, int n_per_row, const block_q8_K * x, block_q8_k_r16 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%16 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_q8_K * x16[16]; + for (int row = 0; row < nrows; row += 16) { + for (int k = 0; k < 16; ++k) x16[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 16; ++k) { + y[ibl].d[k] = GGML_FP32_TO_FP16(x16[k][ibl].d); + for (int ib = 0; ib < QK_K/4; ++ib) { + for (int i = 0; i < 4; ++i) y[ibl].qs[64*ib + 4*k + i] = x16[k][ibl].qs[4*ib+i]; + } + } +#ifdef HAVE_FANCY_SIMD + for (int l = 0; l < 64; ++l) { + auto v = _mm512_xor_si512(_mm512_loadu_si512((const __m512i *)y[ibl].qs + l), _mm512_set1_epi8(-128)); + _mm512_storeu_si512((__m512i *)y[ibl].qs + l, v); + } +#endif + } + x += 16*nblock; + y += nblock; + } +} + +size_t quantize_q8_k_r16(const float * src, void * dst, int64_t nrows, int64_t n_per_row, [[maybe_unused]] const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%16 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + char * qcur = (char *)dst; + auto row_size_0 = ggml_row_size(GGML_TYPE_Q8_K, n_per_row); + auto row_size_1 = ggml_row_size(GGML_TYPE_Q8_K_R16, n_per_row); + std::vector qtmp(16*row_size_0); + for (int row = 0; row < nrows; row += 16) { + quantize_row_q8_K32(src, (void *)qtmp.data(), 16*n_per_row); + repack_q16_k(16, n_per_row, (const block_q8_K *)qtmp.data(), (block_q8_k_r16 *)qcur, false); + qcur += 16*row_size_1; + src += 16*n_per_row; + } + return nrows*row_size_1; +} + +void dequantize_row_q8_k_r16(const block_q8_k_r16 * x, float * y, int64_t k) { + auto n_per_row = k/16; + float * y16[16]; + for (int k = 0; k < 16; ++k) y16[k] = y + n_per_row*k; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto qs = (const uint8_t *)x[ibl].qs; + for (int k = 0; k < 16; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + const float m = -128.f*d; + for (int ib = 0; ib < QK_K/4; ++ib) { + for (int i = 0; i < 4; ++i) { + y16[k][QK_K*ibl+4*ib+i] = d * qs[64*ib+4*k+i] + m; + } + } + } + } +} + +void vec_dot_q8_k_r16_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q8_K_R16, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= q8_KV_r8 +// + +void quantize_row_q8_KV_r8_ref(const float * x, void * y, int64_t k) { + quantize_q8_KV_r8(x, y, 8, k/8, nullptr, nullptr); +} + +void quantize_row_q8_KV_r8(const float * x, void * y, int64_t k) { + quantize_q8_KV_r8(x, y, 8, k/8, nullptr, nullptr); +} + +static void repack_q8_KV(int nrows, int n_per_row, const char * cx, char * cy, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%8 == 0); + GGML_ASSERT(n_per_row%16 == 0); + auto row_size_x = ggml_row_size(GGML_TYPE_Q8_KV, n_per_row); + auto row_size_y = ggml_row_size(GGML_TYPE_Q8_KV_R8, n_per_row); + const int8_t * x8[8]; +#ifdef __ARM_NEON + int8x16x2_t m0, m1, m2, m3; +#endif + for (int row = 0; row < nrows; row += 8) { + auto dy = (float *)cy; + auto qy = (int8_t *)(dy + 8); + for (int k = 0; k < 8; ++k) { + auto dx = (const float *)(cx + k*row_size_x); + dy[k] = dx[0]; + x8[k] = (const int8_t *)(dx + 2); + } + for (int ib = 0; ib < n_per_row/16; ++ib) { +#ifdef __AVX2__ +#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1) + auto m0 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[4]+ib), _mm_loadu_si128((const __m128i *)x8[0]+ib)); + auto m1 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[5]+ib), _mm_loadu_si128((const __m128i *)x8[1]+ib)); + auto m2 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[6]+ib), _mm_loadu_si128((const __m128i *)x8[2]+ib)); + auto m3 = MM256_SET_M128I(_mm_loadu_si128((const __m128i *)x8[7]+ib), _mm_loadu_si128((const __m128i *)x8[3]+ib)); + auto t0 = _mm256_unpacklo_epi32(m0, m1); + auto t1 = _mm256_unpacklo_epi32(m2, m3); + auto t2 = _mm256_unpackhi_epi32(m0, m1); + auto t3 = _mm256_unpackhi_epi32(m2, m3); + m0 = _mm256_unpacklo_epi64(t0, t1); + m1 = _mm256_unpackhi_epi64(t0, t1); + m2 = _mm256_unpacklo_epi64(t2, t3); + m3 = _mm256_unpackhi_epi64(t2, t3); +#ifdef HAVE_FANCY_SIMD + if (online) { + m0 = _mm256_add_epi8(m0, _mm256_set1_epi8(127)); + m1 = _mm256_add_epi8(m1, _mm256_set1_epi8(127)); + m2 = _mm256_add_epi8(m2, _mm256_set1_epi8(127)); + m3 = _mm256_add_epi8(m3, _mm256_set1_epi8(127)); + } +#endif + _mm256_storeu_si256((__m256i *)qy + 4*ib+0, m0); + _mm256_storeu_si256((__m256i *)qy + 4*ib+1, m1); + _mm256_storeu_si256((__m256i *)qy + 4*ib+2, m2); + _mm256_storeu_si256((__m256i *)qy + 4*ib+3, m3); +#elif defined __ARM_NEON + m0.val[0] = vld1q_s8(x8[0]+16*ib); m0.val[1] = vld1q_s8(x8[4]+16*ib); + m1.val[0] = vld1q_s8(x8[1]+16*ib); m1.val[1] = vld1q_s8(x8[5]+16*ib); + m2.val[0] = vld1q_s8(x8[2]+16*ib); m2.val[1] = vld1q_s8(x8[6]+16*ib); + m3.val[0] = vld1q_s8(x8[3]+16*ib); m3.val[1] = vld1q_s8(x8[7]+16*ib); + auto row01 = vtrnq_s32(vreinterpretq_s32_s8(m0.val[0]), vreinterpretq_s32_s8(m1.val[0])); + auto row23 = vtrnq_s32(vreinterpretq_s32_s8(m2.val[0]), vreinterpretq_s32_s8(m3.val[0])); + m0.val[0] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m1.val[0] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + m2.val[0] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m3.val[0] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + row01 = vtrnq_s32(vreinterpretq_s32_s8(m0.val[1]), vreinterpretq_s32_s8(m1.val[1])); + row23 = vtrnq_s32(vreinterpretq_s32_s8(m2.val[1]), vreinterpretq_s32_s8(m3.val[1])); + m0.val[1] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m1.val[1] = vreinterpretq_s8_s64(vtrn1q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + m2.val[1] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[0]), vreinterpretq_s64_s32(row23.val[0]))); + m3.val[1] = vreinterpretq_s8_s64(vtrn2q_s64(vreinterpretq_s64_s32(row01.val[1]), vreinterpretq_s64_s32(row23.val[1]))); + vst1q_s8_x2(qy + 0 + 128*ib, m0); + vst1q_s8_x2(qy + 32 + 128*ib, m1); + vst1q_s8_x2(qy + 64 + 128*ib, m2); + vst1q_s8_x2(qy + 96 + 128*ib, m3); +#else + // TODO + for (int l = 0; l < 4; ++l) { + for (int k = 0; k < 8; ++k) for (int i = 0; i < 4; ++i) { + y[ib].qs[32*l+4*k+i+ 0] = x8[k][ib].qs[i+4*l+ 0]; + y[ib].qs[32*l+4*k+i+128] = x8[k][ib].qs[i+4*l+16]; + } + } +#endif + + } + cx += 8*row_size_x; + cy += online ? 8*row_size_x : 8*row_size_y; + //So, if we are run-time-repacking (online = true) we don't want to change the stride, so we just leave some unused space at the end of each row + } +} +#ifdef HAVE_FANCY_SIMD +static void modify_q8_KV_r8(int64_t k, char * cy) { + int8_t * q8 = (int8_t *)(cy + 8*sizeof(float)); + for (int j = 0; j < k; ++j) q8[j] += 127; +} +#endif + +size_t quantize_q8_KV_r8(const float * src, void * dst, int64_t nrows, int64_t n_per_row, [[maybe_unused]] const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%8 == 0); + GGML_ASSERT(n_per_row%16 == 0); + char * qcur = (char *)dst; + auto row_size_0 = ggml_row_size(GGML_TYPE_Q8_KV, n_per_row); + auto row_size_1 = ggml_row_size(GGML_TYPE_Q8_KV_R8, n_per_row); + std::vector qtmp(8*row_size_0); + for (int row = 0; row < nrows; row += 8) { + quantize_q8_KV(src, (void *)qtmp.data(), 8, n_per_row, imatrix, user_data); + repack_q8_KV(8, n_per_row, qtmp.data(), qcur, false); + qcur += 8*row_size_1; + src += 8*n_per_row; + } + return nrows*row_size_1; +} + +void dequantize_row_q8_KV_r8(const void * vx, float * y, int64_t k) { + auto n_per_row = k/8; + float * y8[8]; + for (int k = 0; k < 8; ++k) y8[k] = y + n_per_row*k; + auto dptr = (const float *)vx; + auto q8 = (const int8_t *)(dptr + 8); + for (int ib = 0; ib < n_per_row/16; ++ib) { + for (int k = 0; k < 8; ++k) { + for (int l = 0; l < 4; ++l) { + for (int i = 0; i < 4; ++i) y8[k][16*ib + 4*l + i] = dptr[k] * q8[128*ib + 32*l + 4*k + i]; + } + } + } +} + +void vec_dot_q8_KV_r8_q8_KV(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q8_KV_R8, vx, 0, GGML_TYPE_Q8_KV, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= bf16_r4 +// +namespace { +inline ggml_bf16_t to_bf16(const float& x) { + union { float f; uint32_t u; } helper; + helper.f = x; + return ggml_bf16_t{(uint16_t)(helper.u >> 16)}; +} +inline ggml_bf16_t to_bf16(const ggml_half& x) { return to_bf16(GGML_FP16_TO_FP32(x)); } +inline ggml_bf16_t to_bf16(const ggml_bf16_t& x) { return x; } +template +void repack_bf16(int nrows, int n_per_row, const T * x, ggml_bf16_t * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%16 == 0); + GGML_ASSERT(n_per_row%2 == 0); + for (int row = 0; row < nrows; row += 16) { + for (int k = 0; k < 16; ++k) { + auto x8 = x + k*n_per_row; + for (int ib = 0; ib < n_per_row/2; ++ib) { + y[32*ib + 2*k + 0] = to_bf16(x8[2*ib+0]); + y[32*ib + 2*k + 1] = to_bf16(x8[2*ib+1]); + } + } + x += 16*n_per_row; + y += 16*n_per_row; + } +} +} + +void repack_f32_bf16_r16(const void * src, void * dst, int64_t nrows, int64_t n_per_row) { + repack_bf16(nrows, n_per_row, (const float *)src, (ggml_bf16_t *)dst, false); +} + +void repack_bf16_bf16_r16(const void * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row) { + repack_bf16(nrows, n_per_row, (const ggml_bf16_t *)src, (ggml_bf16_t *)dst, false); +} + +// +// ========================================= iq3_k_r4 +// + +void quantize_row_iq3_k_r4_ref(const float * x, block_iq3_k_r4 * y, int64_t k) { + quantize_iq3_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq3_k_r4(const float * x, void * y, int64_t k) { + quantize_iq3_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +inline void convert_iq3_k(const block_iq3_k& x, uint8_t * L) { + const uint8_t * qs = x.qs; + const uint8_t * qh = x.qh; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + int shift_l = 2*(ib32%4); + int shift_h = ib32%8; + for (int j = 0; j < 16; ++j) { + L[j+ 0] = ((qs[j+ 0] >> shift_l) & 3) | (((qh[j+ 0] >> shift_h) & 1) << 2); + L[j+16] = ((qs[j+16] >> shift_l) & 3) | (((qh[j+16] >> shift_h) & 1) << 2); + } + L += 32; + if (shift_l == 6) qs += 32; + } +} +} + +static void repack_iq3_k(int nrows, int n_per_row, const block_iq3_k * x, block_iq3_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq3_k * x4[4]; + uint8_t L[QK_K]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].extra, 0, 8); + std::memset(y[ibl].scales_l, 0, QK_K/8); + std::memset(y[ibl].scales_h, 0, QK_K/32); + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + auto extra = x4[k][ibl].extra; + uint16_t sh = x4[k][ibl].scales_h; + convert_iq3_k(x4[k][ibl], L); + for (int ib = 0; ib < QK_K/32; ++ib) { + if (extra & 1) y[ibl].extra[k+0] |= (1 << ib); + if (extra & 2) y[ibl].extra[k+4] |= (1 << ib); + extra >>= 2; + uint8_t sl1 = x4[k][ibl].scales_l[ib] & 0xf; + uint8_t sl2 = x4[k][ibl].scales_l[ib] >> 4; + uint8_t sh1 = (sh >> 0) & 1; + uint8_t sh2 = (sh >> 1) & 1; + sh >>= 2; + int i = 8*ib + k; + y[ibl].scales_l[i%32] |= (sl1 << 4*(i/32)); + y[ibl].scales_h[i%8 ] |= (sh1 << (i/8)); + i += 4; + y[ibl].scales_l[i%32] |= (sl2 << 4*(i/32)); + y[ibl].scales_h[i%8 ] |= (sh2 << (i/8)); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[32*ib+4*k+i+ 0] = ((L[32*ib+i+ 0] & 0x3) << 0) | ((L[32*ib+i+ 4] & 0x3) << 2) | ((L[32*ib+i+ 8] & 0x3) << 4) | ((L[32*ib+i+12] & 0x3) << 6); + y[ibl].qs[32*ib+4*k+i+16] = ((L[32*ib+i+16] & 0x3) << 0) | ((L[32*ib+i+20] & 0x3) << 2) | ((L[32*ib+i+24] & 0x3) << 4) | ((L[32*ib+i+28] & 0x3) << 6); + y[ibl].qh[16*ib+4*k+i+ 0] = ((L[32*ib+i+ 0] >> 2) << 0) | ((L[32*ib+i+ 4] >> 2) << 1) | ((L[32*ib+i+ 8] >> 2) << 2) | ((L[32*ib+i+12] >> 2) << 3) + | ((L[32*ib+i+16] >> 2) << 4) | ((L[32*ib+i+20] >> 2) << 5) | ((L[32*ib+i+24] >> 2) << 6) | ((L[32*ib+i+28] >> 2) << 7); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq3_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + char * qcur = (char *)dst; + auto row_size = ggml_row_size(GGML_TYPE_IQ3_K, n_per_row); + std::vector qtmp(4*row_size); + for (int row = 0; row < nrows; row += 4) { + quantize_iq3_k(src, (void *)qtmp.data(), 4, n_per_row, imatrix, user_data); + repack_iq3_k(4, n_per_row, (const block_iq3_k *)qtmp.data(), (block_iq3_k_r4 *)qcur, false); + qcur += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq3_k_r4(const block_iq3_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + auto ql = x[ibl].qs; + auto qh = x[ibl].qh; + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 8*ib + k; + float dl1 = d * (2*((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) + 1) * ((x[ibl].scales_h[is%8] >> (is/8)) & 1 ? -1 : 1); + is += 4; + float dl2 = d * (2*((x[ibl].scales_l[is%32] >> 4*(is/32)) & 0xf) + 1) * ((x[ibl].scales_h[is%8] >> (is/8)) & 1 ? -1 : 1); + auto values1 = iq3nl_values + (x[ibl].extra[k+0] & (1 << ib) ? 8 : 0); + auto values2 = iq3nl_values + (x[ibl].extra[k+4] & (1 << ib) ? 8 : 0); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl1 * values1[((ql[4*k+i+ 0] >> 0) & 3) | ((qh[4*k+i] << 2) & 4)]; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl1 * values1[((ql[4*k+i+ 0] >> 2) & 3) | ((qh[4*k+i] << 1) & 4)]; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl1 * values1[((ql[4*k+i+ 0] >> 4) & 3) | ((qh[4*k+i] << 0) & 4)]; + y4[k][QK_K*ibl+32*ib+i+12] = dl1 * values1[((ql[4*k+i+ 0] >> 6) & 3) | ((qh[4*k+i] >> 1) & 4)]; + y4[k][QK_K*ibl+32*ib+i+16] = dl2 * values2[((ql[4*k+i+16] >> 0) & 3) | ((qh[4*k+i] >> 2) & 4)]; + y4[k][QK_K*ibl+32*ib+i+20] = dl2 * values2[((ql[4*k+i+16] >> 2) & 3) | ((qh[4*k+i] >> 3) & 4)]; + y4[k][QK_K*ibl+32*ib+i+24] = dl2 * values2[((ql[4*k+i+16] >> 4) & 3) | ((qh[4*k+i] >> 4) & 4)]; + y4[k][QK_K*ibl+32*ib+i+28] = dl2 * values2[((ql[4*k+i+16] >> 6) & 3) | ((qh[4*k+i] >> 5) & 4)]; + } + ql += 32; + qh += 16; + } + } + } +} + +void vec_dot_iq3_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ3_K_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq2_k_r4 +// + +void quantize_row_iq2_k_r4_ref(const float * x, block_iq2_k_r4 * y, int64_t k) { + quantize_iq2_k_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq2_k_r4(const float * x, void * y, int64_t k) { + quantize_iq2_k_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +inline void convert_iq2_k(const block_iq2_k& x, uint8_t * L) { + const uint8_t * qs = x.qs; + for (int ib32 = 0; ib32 < QK_K/32; ++ib32) { + int shift_l = 2*(ib32%4); + for (int j = 0; j < 16; ++j) { + L[j+ 0] = ((qs[j+ 0] >> shift_l) & 3); + L[j+16] = ((qs[j+16] >> shift_l) & 3); + } + L += 32; + if (shift_l == 6) qs += 32; + } +} +} + +static void repack_iq2_k(int nrows, int n_per_row, const block_iq2_k * x, block_iq2_k_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq2_k * x4[4]; + uint8_t L[QK_K]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].extra, 0, 8); + std::memset(y[ibl].scales, 0, QK_K/8); + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + auto extra = x4[k][ibl].extra; + convert_iq2_k(x4[k][ibl], L); + for (int ib = 0; ib < QK_K/32; ++ib) { + if (extra & 1) y[ibl].extra[k+0] |= (1 << ib); + if (extra & 2) y[ibl].extra[k+4] |= (1 << ib); + extra >>= 2; + uint8_t sl1 = x4[k][ibl].scales[ib] & 0xf; + uint8_t sl2 = x4[k][ibl].scales[ib] >> 4; + int i = 8*ib + k; + y[ibl].scales[i%32] |= (sl1 << 4*(i/32)); + i += 4; + y[ibl].scales[i%32] |= (sl2 << 4*(i/32)); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[32*ib+4*k+i+ 0] = ((L[32*ib+i+ 0] & 0x3) << 0) | ((L[32*ib+i+ 4] & 0x3) << 2) | ((L[32*ib+i+ 8] & 0x3) << 4) | ((L[32*ib+i+12] & 0x3) << 6); + y[ibl].qs[32*ib+4*k+i+16] = ((L[32*ib+i+16] & 0x3) << 0) | ((L[32*ib+i+20] & 0x3) << 2) | ((L[32*ib+i+24] & 0x3) << 4) | ((L[32*ib+i+28] & 0x3) << 6); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq2_k_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + char * qcur = (char *)dst; + auto row_size = ggml_row_size(GGML_TYPE_IQ2_K, n_per_row); + std::vector qtmp(4*row_size); + for (int row = 0; row < nrows; row += 4) { + quantize_iq2_k(src, (void *)qtmp.data(), 4, n_per_row, imatrix, user_data); + repack_iq2_k(4, n_per_row, (const block_iq2_k *)qtmp.data(), (block_iq2_k_r4 *)qcur, false); + qcur += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq2_k_r4(const block_iq2_k_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + auto ql = x[ibl].qs; + for (int ib = 0; ib < QK_K/32; ++ib) { + int is = 8*ib + k; + float dl1 = d * (((x[ibl].scales[is%32] >> 4*(is/32)) & 0xf) - 8); + is += 4; + float dl2 = d * (((x[ibl].scales[is%32] >> 4*(is/32)) & 0xf) - 8); + auto values1 = iq2nl_values + (x[ibl].extra[k+0] & (1 << ib) ? 4 : 0); + auto values2 = iq2nl_values + (x[ibl].extra[k+4] & (1 << ib) ? 4 : 0); + for (int i = 0; i < 4; ++i) { + y4[k][QK_K*ibl+32*ib+i+ 0] = dl1 * values1[(ql[4*k+i+ 0] >> 0) & 3]; + y4[k][QK_K*ibl+32*ib+i+ 4] = dl1 * values1[(ql[4*k+i+ 0] >> 2) & 3]; + y4[k][QK_K*ibl+32*ib+i+ 8] = dl1 * values1[(ql[4*k+i+ 0] >> 4) & 3]; + y4[k][QK_K*ibl+32*ib+i+12] = dl1 * values1[(ql[4*k+i+ 0] >> 6) & 3]; + y4[k][QK_K*ibl+32*ib+i+16] = dl2 * values2[(ql[4*k+i+16] >> 0) & 3]; + y4[k][QK_K*ibl+32*ib+i+20] = dl2 * values2[(ql[4*k+i+16] >> 2) & 3]; + y4[k][QK_K*ibl+32*ib+i+24] = dl2 * values2[(ql[4*k+i+16] >> 4) & 3]; + y4[k][QK_K*ibl+32*ib+i+28] = dl2 * values2[(ql[4*k+i+16] >> 6) & 3]; + } + ql += 32; + } + } + } +} + +void vec_dot_iq2_k_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_K_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +namespace { +inline uint8_t scrambled_sign(uint8_t s) { + static const uint8_t k_table[128] = { + 0x00, 0x7f, 0x7e, 0x01, 0x7c, 0x03, 0x02, 0x7d, 0x78, 0x07, 0x06, 0x79, 0x04, 0x7b, 0x7a, 0x05, + 0x70, 0x0f, 0x0e, 0x71, 0x0c, 0x73, 0x72, 0x0d, 0x08, 0x77, 0x76, 0x09, 0x74, 0x0b, 0x0a, 0x75, + 0x60, 0x1f, 0x1e, 0x61, 0x1c, 0x63, 0x62, 0x1d, 0x18, 0x67, 0x66, 0x19, 0x64, 0x1b, 0x1a, 0x65, + 0x10, 0x6f, 0x6e, 0x11, 0x6c, 0x13, 0x12, 0x6d, 0x68, 0x17, 0x16, 0x69, 0x14, 0x6b, 0x6a, 0x15, + 0x40, 0x3f, 0x3e, 0x41, 0x3c, 0x43, 0x42, 0x3d, 0x38, 0x47, 0x46, 0x39, 0x44, 0x3b, 0x3a, 0x45, + 0x30, 0x4f, 0x4e, 0x31, 0x4c, 0x33, 0x32, 0x4d, 0x48, 0x37, 0x36, 0x49, 0x34, 0x4b, 0x4a, 0x35, + 0x20, 0x5f, 0x5e, 0x21, 0x5c, 0x23, 0x22, 0x5d, 0x58, 0x27, 0x26, 0x59, 0x24, 0x5b, 0x5a, 0x25, + 0x50, 0x2f, 0x2e, 0x51, 0x2c, 0x53, 0x52, 0x2d, 0x28, 0x57, 0x56, 0x29, 0x54, 0x2b, 0x2a, 0x55, + }; + return k_table[s]; +} +} + +// +// ========================================= iq2_xxs_r4 +// + +void quantize_row_iq2_xxs_r4_ref(const float * x, block_iq2_xxs_r4 * y, int64_t k) { + quantize_iq2_xxs_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq2_xxs_r4(const float * x, void * y, int64_t k) { + quantize_iq2_xxs_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static void repack_iq2_xxs(int nrows, int n_per_row, const block_iq2_xxs * x, block_iq2_xxs_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq2_xxs * x4[4]; + uint32_t aux32[2]; + const uint8_t * aux8 = (const uint8_t *)aux32; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto ysas = (uint32_t *)y[ibl].sas; + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + for (int ib = 0; ib < QK_K/32; ++ib) { + std::memcpy(aux32, x4[k][ibl].qs + 4*ib, 2*sizeof(uint32_t)); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[16*ib+4*k+i] = aux8[i]; + } + uint8_t scale = aux32[1] >> 28; + uint8_t s1 = (scrambled_sign((aux32[1] >> 0) & 127) << 1) | ((scale >> 0) & 1); + uint8_t s2 = (scrambled_sign((aux32[1] >> 7) & 127) << 1) | ((scale >> 1) & 1); + uint8_t s3 = (scrambled_sign((aux32[1] >> 14) & 127) << 1) | ((scale >> 2) & 1); + uint8_t s4 = (scrambled_sign((aux32[1] >> 21) & 127) << 1) | ((scale >> 3) & 1); + aux32[1] = uint32_t(s1) | (uint32_t(s2) << 8) | (uint32_t(s3) << 16) | (uint32_t(s4) << 24); + ysas[4*ib+k] = aux32[1]; + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq2_xxs_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_iq2_xxs(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<32, block_iq2_xxs, block_iq2_xxs_r4, 4>(GGML_TYPE_IQ2_XXS, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_iq2_xxs); +} + +void dequantize_row_iq2_xxs_r4(const block_iq2_xxs_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + uint32_t s32; + const uint8_t * s8 = (const uint8_t *)&s32; + for (int ibl = 0; ibl < nblock; ++ibl) { + const uint32_t * sas = (const uint32_t *)x[ibl].sas; + for (int k = 0; k < 4; ++k) { + const float d = 0.125f*GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/32; ++ib) { + uint32_t aux32 = sas[4*ib+k]; + s32 = aux32 & 0x01010101; + uint8_t scale = s8[0] | (s8[1] << 1) | (s8[2] << 2) | (s8[3] << 3); + float dl = d*(2*scale+1); + aux32 &= 0xfefefefe; + aux32 ^= (aux32 >> 1); + for (int i = 0; i < 4; ++i) { + auto val = (const int8_t *)(iq2xxs_grid + x[ibl].qs[16*ib+4*k+i]); + for (int j = 0; j < 8; ++j) y4[k][QK_K*ibl+32*ib+8*i+j] = dl * val[j] * (aux32 & (1 << j) ? -1 : 1); + aux32 >>= 8; + } + } + } + } +} + +void vec_dot_iq2_xxs_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_XXS_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq2_xs_r4 +// + +void quantize_row_iq2_xs_r4_ref(const float * x, block_iq2_xs_r4 * y, int64_t k) { + quantize_iq2_xs_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq2_xs_r4(const float * x, void * y, int64_t k) { + quantize_iq2_xs_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static void repack_iq2_xs(int nrows, int n_per_row, const block_iq2_xs * x, block_iq2_xs_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq2_xs * x4[4]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + for (int ib = 0; ib < QK_K/32; ++ib) { + for (int i = 0; i < 4; ++i) { + uint16_t v = x4[k][ibl].qs[4*ib+i]; + uint8_t s = v >> 9; + y[ibl].qs[16*ib+4*k+i] = (v & 511) | (scrambled_sign(s) << 9); + } + y[ibl].scales[4*ib+k] = x4[k][ibl].scales[ib]; + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq2_xs_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_iq2_xs(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<16, block_iq2_xs, block_iq2_xs_r4, 4>(GGML_TYPE_IQ2_XS, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_iq2_xs); +} + +void dequantize_row_iq2_xs_r4(const block_iq2_xs_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = 0.125f*GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/32; ++ib) { + float dl1 = d * (2*(x[ibl].scales[4*ib+k] & 0xf) + 1); + float dl2 = d * (2*(x[ibl].scales[4*ib+k] >> 4) + 1); + for (int i = 0; i < 4; ++i) { + auto val = (const int8_t *)(iq2xs_grid + (x[ibl].qs[16*ib+4*k+i] & 511)); + auto signs = x[ibl].qs[16*ib+4*k+i] >> 9; + signs ^= (signs << 1); + float dl = i < 2 ? dl1 : dl2; + for (int j = 0; j < 8; ++j) y4[k][QK_K*ibl+32*ib+8*i+j] = dl * val[j] * (signs & (1 << j) ? -1 : 1); + } + } + } + } +} + +void vec_dot_iq2_xs_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_XS_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq2_s_r4 +// + +void quantize_row_iq2_s_r4_ref(const float * x, block_iq2_s_r4 * y, int64_t k) { + quantize_iq2_s_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq2_s_r4(const float * x, void * y, int64_t k) { + quantize_iq2_s_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static void repack_iq2_s(int nrows, int n_per_row, const block_iq2_s * x, block_iq2_s_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq2_s * x4[4]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + auto signs = x4[k][ibl].qs + QK_K/8; + y[ibl].d[k] = x4[k][ibl].d; + for (int ib = 0; ib < QK_K/32; ++ib) { + y[ibl].scales[4*ib+k] = x4[k][ibl].scales[ib]; + for (int i = 0; i < 4; ++i) { + y[ibl].qs[16*ib+4*k+i] = x4[k][ibl].qs[4*ib+i]; + y[ibl].signs[16*ib+4*k+i] = signs[4*ib+i]; + } + y[ibl].qh[4*ib+k] = x4[k][ibl].qh[ib]; + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq2_s_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_iq2_s(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<16, block_iq2_s, block_iq2_s_r4, 4>(GGML_TYPE_IQ2_S, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_iq2_s); +} + +void dequantize_row_iq2_s_r4(const block_iq2_s_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = 0.125f*GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/32; ++ib) { + float dl1 = d * (2*(x[ibl].scales[4*ib+k] & 0xf) + 1); + float dl2 = d * (2*(x[ibl].scales[4*ib+k] >> 4) + 1); + for (int i = 0; i < 4; ++i) { + auto val = (const int8_t *)(iq2s_grid + (x[ibl].qs[16*ib+4*k+i] | ((x[ibl].qh[4*ib+k] << (8 - 2*i)) & 0x300))); + auto signs = x[ibl].signs[16*ib+4*k+i]; + float dl = i < 2 ? dl1 : dl2; + for (int j = 0; j < 8; ++j) y4[k][QK_K*ibl+32*ib+8*i+j] = dl * val[j] * (signs & (1 << j) ? -1 : 1); + } + } + } + } +} + +void vec_dot_iq2_s_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_S_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq3_xxs_r4 +// + +void quantize_row_iq3_xxs_r4_ref(const float * x, block_iq3_xxs_r4 * y, int64_t k) { + quantize_iq3_xxs_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq3_xxs_r4(const float * x, void * y, int64_t k) { + quantize_iq3_xxs_r4(x, y, 4, k/4, nullptr, nullptr); +} + +namespace { +} + +static void repack_iq3_xxs(int nrows, int n_per_row, const block_iq3_xxs * x, block_iq3_xxs_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq3_xxs * x4[4]; + uint32_t aux32; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto ysas = (uint32_t *)y[ibl].sas; + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + auto xsas = x4[k][ibl].qs + QK_K/4; + for (int ib = 0; ib < QK_K/32; ++ib) { + for (int i = 0; i < 8; ++i) { + y[ibl].qs[32*ib+8*k+i] = x4[k][ibl].qs[8*ib+i]; + } + std::memcpy(&aux32, xsas + 4*ib, 4); + uint8_t scale = aux32 >> 28; + uint8_t s1 = (scrambled_sign((aux32 >> 0) & 127) << 1) | ((scale >> 0) & 1); + uint8_t s2 = (scrambled_sign((aux32 >> 7) & 127) << 1) | ((scale >> 1) & 1); + uint8_t s3 = (scrambled_sign((aux32 >> 14) & 127) << 1) | ((scale >> 2) & 1); + uint8_t s4 = (scrambled_sign((aux32 >> 21) & 127) << 1) | ((scale >> 3) & 1); + aux32 = uint32_t(s1) | (uint32_t(s2) << 8) | (uint32_t(s3) << 16) | (uint32_t(s4) << 24); + ysas[4*ib+k] = aux32; + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq3_xxs_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_iq3_xxs(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<32, block_iq3_xxs, block_iq3_xxs_r4, 4>(GGML_TYPE_IQ3_XXS, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_iq3_xxs); +} + +void dequantize_row_iq3_xxs_r4(const block_iq3_xxs_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + uint32_t s32; + const uint8_t * s8 = (const uint8_t *)&s32; + for (int ibl = 0; ibl < nblock; ++ibl) { + const uint32_t * sas = (const uint32_t *)x[ibl].sas; + for (int k = 0; k < 4; ++k) { + const float d = 0.25f*GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/32; ++ib) { + uint32_t aux32 = sas[4*ib+k]; + s32 = aux32 & 0x01010101; + uint8_t scale = s8[0] | (s8[1] << 1) | (s8[2] << 2) | (s8[3] << 3); + float dl = d*(2*scale+1); + aux32 &= 0xfefefefe; + aux32 ^= (aux32 >> 1); + for (int i = 0; i < 8; ++i) { + auto val = (const int8_t *)(iq3xxs_grid + x[ibl].qs[32*ib+8*k+i]); + for (int j = 0; j < 4; ++j) y4[k][QK_K*ibl+32*ib+4*i+j] = dl * val[j] * (aux32 & (1 << j) ? -1 : 1); + aux32 >>= 4; + } + } + } + } +} + +void vec_dot_iq3_xxs_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ3_XXS_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +// +// ========================================= iq3_s_r4 +// + +void quantize_row_iq3_s_r4_ref(const float * x, block_iq3_s_r4 * y, int64_t k) { + quantize_iq3_s_r4(x, (void *)y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq3_s_r4(const float * x, void * y, int64_t k) { + quantize_iq3_s_r4(x, y, 4, k/4, nullptr, nullptr); +} + +static void repack_iq3_s(int nrows, int n_per_row, const block_iq3_s * x, block_iq3_s_r4 * y, [[maybe_unused]] bool online) { + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%QK_K == 0); + int nblock = n_per_row/QK_K; + const block_iq3_s * x4[4]; + for (int row = 0; row < nrows; row += 4) { + for (int k = 0; k < 4; ++k) x4[k] = x + nblock*k; + for (int ibl = 0; ibl < nblock; ++ibl) { + std::memset(y[ibl].scales, 0, QK_K/16); + std::memset(y[ibl].signs, 0, QK_K/2); + std::memset(y[ibl].qh, 0, QK_K/8); + for (int k = 0; k < 4; ++k) { + y[ibl].d[k] = x4[k][ibl].d; + for (int ib = 0; ib < QK_K/64; ++ib) { + int j = 8*ib + k; + y[ibl].scales[(j+0)%16] |= ((x4[k][ibl].scales[ib] & 0xf) << 4*((j+0)/16)); + y[ibl].scales[(j+4)%16] |= ((x4[k][ibl].scales[ib] >> 4) << 4*((j+4)/16)); + } + for (int ib = 0; ib < QK_K/32; ++ib) { + y[ibl].qh[4*ib+k] = x4[k][ibl].qh[ib]; // leave ot like this? + for (int i = 0; i < 4; ++i) { + y[ibl].qs[32*ib+k+8*i+0] = x4[k][ibl].qs[8*ib+i+0]; + y[ibl].qs[32*ib+k+8*i+4] = x4[k][ibl].qs[8*ib+i+4]; + } + for (int i = 0; i < 4; ++i) { + y[ibl].signs[16*ib+4*k+i] = (((x4[k][ibl].signs[4*ib+0] >> i) & 1) << 0) | (((x4[k][ibl].signs[4*ib+0] >> (4+i)) & 1) << 1) | + (((x4[k][ibl].signs[4*ib+1] >> i) & 1) << 2) | (((x4[k][ibl].signs[4*ib+1] >> (4+i)) & 1) << 3) | + (((x4[k][ibl].signs[4*ib+2] >> i) & 1) << 4) | (((x4[k][ibl].signs[4*ib+2] >> (4+i)) & 1) << 5) | + (((x4[k][ibl].signs[4*ib+3] >> i) & 1) << 6) | (((x4[k][ibl].signs[4*ib+3] >> (4+i)) & 1) << 7); + } + } + } + } + x += 4*nblock; + y += nblock; + } +} + +size_t quantize_iq3_s_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + auto q_func = [] (const float * x, void * vy, int n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + iqkbase_quantize_iq3_s(x, (char *)vy, 1, n_per_row, imatrix, user_data); + }; + return quantize_repack<16, block_iq3_s, block_iq3_s_r4, 4>(GGML_TYPE_IQ3_S, src, dst, nrows, n_per_row, imatrix, user_data, + q_func, repack_iq3_s); +} + +void dequantize_row_iq3_s_r4(const block_iq3_s_r4 * x, float * y, int64_t k) { + auto n_per_row = k/4; + float * y4[4] = {y, y + n_per_row, y + 2*n_per_row, y + 3*n_per_row}; + int nblock = n_per_row/QK_K; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + const float d = GGML_FP16_TO_FP32(x[ibl].d[k]); + for (int ib = 0; ib < QK_K/32; ++ib) { + int l = 4*ib + k; + float dl = d * (1 + 2*((x[ibl].scales[l%16] >> 4*(l/16)) & 0xf)); + for (int i = 0; i < 4; ++i) { + auto grid1 = (const uint8_t *)(iq3s_grid + x[ibl].qs[32*ib+k+8*i+0] + ((x[ibl].qh[4*ib+k] << (8-i)) & 0x100)); + auto grid2 = (const uint8_t *)(iq3s_grid + x[ibl].qs[32*ib+k+8*i+4] + ((x[ibl].qh[4*ib+k] << (4-i)) & 0x100)); + for (int j = 0; j < 4; ++j) { + y4[k][QK_K*ibl+32*ib+4*i+ 0+j] = dl * grid1[j] * (x[ibl].signs[16*ib+4*k+j] & (1 << (i+0)) ? -1 : 1); + y4[k][QK_K*ibl+32*ib+4*i+16+j] = dl * grid2[j] * (x[ibl].signs[16*ib+4*k+j] & (1 << (i+4)) ? -1 : 1); + } + } + } + } + } +} + +void vec_dot_iq3_s_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ3_S_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +void quantize_row_iq1_s_r4_ref(const float * x, block_iq1_s_r4 * y, int64_t k) { + quantize_iq1_s_r4(x, y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq1_s_r4(const float * x, void * y, int64_t k) { + quantize_iq1_s_r4(x, y, 4, k/4, nullptr, nullptr); +} + +size_t quantize_iq1_s_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + constexpr int kBlockSize = 32; + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%kBlockSize == 0); + int nblock = n_per_row/kBlockSize; + float weight[kBlockSize]; + int8_t L[kBlockSize]; + float pairs[2*kBlockSize]; + float sumx[kBlockSize+1], sumw[kBlockSize+1]; + float max[4]; + uint16_t index[4]; + int shift; + float invd[4]; + std::vector scales(4*nblock); + auto row_size = ggml_row_size(GGML_TYPE_IQ1_S_R4, n_per_row); + char * cy = (char *)dst; + for (int row = 0; row < nrows; row += 4) { + ggml_half * dptr = (ggml_half *)cy; + auto y = (block_iq1_s_r4 *)(dptr + 4); + for (int k = 0; k < 4; ++k) max[k] = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + auto xb = src + k*n_per_row + kBlockSize*ibl; + float sumx2 = 0; + for (int j = 0; j < kBlockSize; ++j) sumx2 += xb[j]*xb[j]; + if (sumx2 < 1e-14f) { + //printf("Found block with all zeros\n"); + // all zero + int ind = 1029; // this is the grid entry with all zeros + scales[4*ibl+k] = 0; + uint16_t h = 0; + for (int i = 0; i < 4; ++i) { + y[ibl].qs[4*i + k] = ind & 255; + h |= (ind >> 8) << 3*i; + } + y[ibl].qh[k] = h; + continue; + } + float sigma2 = 1.5f*sumx2/kBlockSize; + bool have_imatrix = false; + if (imatrix) { + have_imatrix = true; + float sumwx = 0; + for (int j = 0; j < kBlockSize; ++j) { + weight[j] = imatrix[kBlockSize*ibl + j]*sqrt(sigma2 + xb[j]*xb[j]); + sumwx += weight[j]*std::abs(xb[j]); + } + if (sumwx < 1e-14f) { + printf("Found block with mismatching importance/model weights\n"); + // Either all weights are zero, or xb is zero where weight is not zero. + // In both of these cases it is better to simply ignore the imatrix + have_imatrix = false; + } + } + if (!have_imatrix) { + for (int j = 0; j < kBlockSize; ++j) weight[j] = sqrt(sigma2 + xb[j]*xb[j]); + } + iqkbase_iq1s_process_1block(kBlockSize, xb, weight, L, scales.data() + 4*ibl + k, index, &shift, pairs, sumx, sumw); + GGML_ASSERT(scales[4*ibl+k] >= 0); + max[k] = std::max(max[k], scales[4*ibl+k]); + uint16_t h = 0; + for (int i = 0; i < 4; ++i) { + GGML_ASSERT(index[i] >= 0 && index[i] < 2048); + y[ibl].qs[4*i + k] = index[i] & 255; + h |= (index[i] >> 8) << 3*i; + } + if (shift < 0) h |= 0x8000; + y[ibl].qh[k] = h; + } + } + for (int k = 0; k < 4; ++k) { + dptr[k] = GGML_FP32_TO_FP16(1.0625f*max[k]/15);; + invd[k] = max[k] ? 15/max[k] : 0.f; + } + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + int ls = nearest_int(0.5f*(scales[4*ibl+k]*invd[k] - 1)); + ls = std::max(0, std::min(7, ls)); + y[ibl].qh[k] |= (ls << 12); + } + } + cy += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq1_s_r4(const block_iq1_s_r4 * x, float * y, int64_t n) { + auto dptr = (const ggml_half *)x; + x = (const block_iq1_s_r4 *)(dptr + 4); + float d[4]; + for (int k = 0; k < 4; ++k) d[k] = GGML_FP16_TO_FP32(dptr[k]); + int n_per_row = n/4; + GGML_ASSERT(n_per_row%32 == 0); + int nblock = n_per_row/32; + float * yk[4]; + for (int k = 0; k < 4; ++k) yk[k] = y + k*n_per_row; + for (int ib = 0; ib < nblock; ++ib) { + for (int k = 0; k < 4; ++k) { + float shift = x[ib].qh[k] & 0x8000 ? -IQ1S_DELTA : IQ1S_DELTA; + float dl = d[k]*(2*((x[ib].qh[k] >> 12) & 7) + 1); + for (int i = 0; i < 4; ++i) { + auto idx = x[ib].qs[4*i+k] | (((x[ib].qh[k] >> 3*i) & 7) << 8); + auto grid = (const int8_t *)(iq1s_grid + idx); + for (int j = 0; j < 8; ++j) yk[k][32*ib + 8*i + j] = dl*(grid[j] + shift); + } + } + } +} + +void vec_dot_iq1_s_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ1_S_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +void quantize_row_iq1_m_r4_ref(const float * x, block_iq1_m_r4 * y, int64_t k) { + quantize_iq1_m_r4(x, y, 4, k/4, nullptr, nullptr); +} + +void quantize_row_iq1_m_r4(const float * x, void * y, int64_t k) { + quantize_iq1_m_r4(x, y, 4, k/4, nullptr, nullptr); +} + +size_t quantize_iq1_m_r4(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + constexpr int kBlockSize = 32; + GGML_ASSERT(nrows%4 == 0); + GGML_ASSERT(n_per_row%kBlockSize == 0); + int nblock = n_per_row/kBlockSize; + float weight[kBlockSize]; + int8_t L[kBlockSize]; + float pairs[2*kBlockSize]; + float max[4]; + uint16_t index[4]; + int shift1, shift2; + float invd[4]; + const uint8_t masks[4] = {0x00, 0x80, 0x08, 0x88}; + std::vector scales(8*nblock); + auto row_size = ggml_row_size(GGML_TYPE_IQ1_M_R4, n_per_row); + char * cy = (char *)dst; + for (int row = 0; row < nrows; row += 4) { + ggml_half * dptr = (ggml_half *)cy; + auto y = (block_iq1_m_r4 *)(dptr + 4); + for (int k = 0; k < 4; ++k) max[k] = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + auto xb = src + k*n_per_row + kBlockSize*ibl; + float sumx2l = 0, sumx2h = 0; + for (int j = 0; j < kBlockSize/2; ++j) sumx2l += xb[j]*xb[j]; + for (int j = kBlockSize/2; j < kBlockSize; ++j) sumx2h += xb[j]*xb[j]; + float sumx2 = sumx2l + sumx2h; + if (sumx2 < 1e-14f) { + scales[8*ibl+2*k+0] = scales[8*ibl+2*k+1] = 0; + int ind = 1029; + for (int i = 0; i < 4; ++i) { + y[ibl].qs[4*i + k] = ind & 255; + } + for (int i = 0; i < 2; ++i) { + y[ibl].qh[4*i+k] = (ind >> 8) | ((ind >> 8) << 4); + } + continue; + } + float sigma2 = 1.5f*sumx2/kBlockSize; + if (imatrix) { + for (int j = 0; j < kBlockSize; ++j) weight[j] = imatrix[kBlockSize*ibl + j]*sqrt(sigma2 + xb[j]*xb[j]); + float sumwx = 0; + for (int j = 0; j < kBlockSize/2; ++j) sumwx += weight[j]*std::abs(xb[j]); + if (sumwx < 1e-14f) { + for (int j = 0; j < kBlockSize/2; ++j) weight[j] = sqrt(sigma2 + xb[j]*xb[j]); + } + sumwx = 0; + for (int j = kBlockSize/2; j < kBlockSize; ++j) sumwx += weight[j]*std::abs(xb[j]); + if (sumwx < 1e-14) { + for (int j = kBlockSize/2; j < kBlockSize; ++j) weight[j] = sqrt(sigma2 + xb[j]*xb[j]); + } + } else { + for (int j = 0; j < kBlockSize; ++j) weight[j] = sqrt(sigma2 + xb[j]*xb[j]); + } + if (sumx2l > 1e-14f) { + iqkbase_iq1m_process_1block(xb+ 0, weight+ 0, L, scales.data() + 8*ibl + 2*k+0, index+0, &shift1, pairs); + } else { + scales[8*ibl+2*k+0] = 0; + index[0] = index[1] = 1029; + } + if (sumx2h > 1e-14f) { + iqkbase_iq1m_process_1block(xb+16, weight+16, L, scales.data() + 8*ibl + 2*k+1, index+2, &shift2, pairs); + } else { + scales[8*ibl+2*k+1] = 0; + index[2] = index[3] = 1029; + } + max[k] = std::max(max[k], std::max(scales[8*ibl+2*k+0], scales[8*ibl+2*k+1])); + for (int i = 0; i < 4; ++i) { + y[ibl].qs[4*i + k] = index[i] & 255; + } + for (int i = 0; i < 2; ++i) { + y[ibl].qh[4*i+k] = (index[2*i+0] >> 8) | ((index[2*i+1] >> 8) << 4); + } + y[ibl].qh[0+k] |= masks[shift1]; + y[ibl].qh[4+k] |= masks[shift2]; + } + } + for (int k = 0; k < 4; ++k) { + dptr[k] = GGML_FP32_TO_FP16(1.0625f*max[k]/15);; + invd[k] = max[k] ? 15/max[k] : 0.f; + } + for (int ibl = 0; ibl < nblock; ++ibl) { + for (int k = 0; k < 4; ++k) { + int ls1 = nearest_int(scales[8*ibl+2*k+0]*invd[k]); + int ls2 = nearest_int(scales[8*ibl+2*k+1]*invd[k]); + ls1 = std::max(0, std::min(15, ls1)); + ls2 = std::max(0, std::min(15, ls2)); + y[ibl].scales[k] = ls1 | (ls2 << 4); + } + } + cy += 4*row_size; + src += 4*n_per_row; + } + return nrows*row_size; +} + +void dequantize_row_iq1_m_r4(const block_iq1_m_r4 * x, float * y, int64_t n) { + auto dptr = (const ggml_half *)x; + x = (const block_iq1_m_r4 *)(dptr + 4); + float d[4]; + for (int k = 0; k < 4; ++k) d[k] = GGML_FP16_TO_FP32(dptr[k]); + int n_per_row = n/4; + GGML_ASSERT(n_per_row%32 == 0); + int nblock = n_per_row/32; + float dl[2]; + float * yk[4]; + for (int k = 0; k < 4; ++k) yk[k] = y + k*n_per_row; + for (int ib = 0; ib < nblock; ++ib) { + for (int k = 0; k < 4; ++k) { + dl[0] = d[k]*(x[ib].scales[k] & 0xf); + dl[1] = d[k]*(x[ib].scales[k] >> 4); + for (int i = 0; i < 2; ++i) { + auto idx1 = x[ib].qs[8*i+k+0] | ((x[ib].qh[4*i+k] & 0x07) << 8); + auto idx2 = x[ib].qs[8*i+k+4] | ((x[ib].qh[4*i+k] & 0x70) << 4); + auto grid1 = (const int8_t *)(iq1s_grid + idx1); + auto grid2 = (const int8_t *)(iq1s_grid + idx2); + auto delta1 = x[ib].qh[4*i+k] & 0x08 ? -IQ1M_DELTA : IQ1M_DELTA; + auto delta2 = x[ib].qh[4*i+k] & 0x80 ? -IQ1M_DELTA : IQ1M_DELTA; + for (int j = 0; j < 8; ++j) yk[k][32*ib + 16*i + j + 0] = dl[i]*(grid1[j] + delta1); + for (int j = 0; j < 8; ++j) yk[k][32*ib + 16*i + j + 8] = dl[i]*(grid2[j] + delta2); + } + } + } +} + +void vec_dot_iq1_m_r4_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ1_M_R4, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + +void quantize_row_q8_KV(const float * x, void * vy, int64_t k) { + iqk_quantize_row_q8_KV(x, vy, k); +} + +void quantize_row_q8_KV_ref(const float * x, void * y, int64_t k) { + quantize_row_q8_KV(x, y, k); +} + +size_t quantize_q8_KV(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + (void)imatrix; + auto row_size = ggml_row_size(GGML_TYPE_Q8_KV, n_per_row); + auto q = (char *)dst; + for (int row = 0; row < nrows; ++row) { + quantize_row_q8_KV(src, q, n_per_row); + src += n_per_row; + q += row_size; + } + return row_size*nrows; +} + +void dequantize_row_q8_KV(const void * x, float * y, int64_t k) { + auto dptr = (const float *)x; + float d = dptr[0]; + auto q8 = (const int8_t *)(dptr + 2); + for (int j = 0; j < k; ++j) y[j] = d * q8[j]; +} + +void vec_dot_q8_KV_q8_KV(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_Q8_KV, vx, 0, GGML_TYPE_Q8_KV, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + GGML_ASSERT(n%QK4_NL == 0); + GGML_ASSERT(nrc == 1); + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); +} + + +//================================================ + +namespace { +struct Repack { + using repack_func = void (*) (int nrows, int n_per_row, const char * src, char * dst, bool online); + ggml_type new_type; + int num_rows; + repack_func repack; +}; +struct Modify { + using modify_func_t = void (*)(int64_t k, char * src_dst); + modify_func_t mod_func; + int nrows; +}; +const Modify * get_modify_info(ggml_type type) { + static const std::unordered_map k_mod_map = { +#ifdef __ARM_NEON + { GGML_TYPE_Q4_0_R8, {modify_q4_0_r8, 8} }, +#endif +#ifdef HAVE_FANCY_SIMD + { GGML_TYPE_Q8_0_R8, {modify_q8_0_r8, 8} }, + { GGML_TYPE_Q8_K_R8, {modify_q8_k_r8, 8} }, + { GGML_TYPE_Q8_KV_R8, {modify_q8_KV_r8, 8} }, +#endif + }; + auto it = k_mod_map.find(type); + return it != k_mod_map.end() ? &it->second : nullptr; +} +bool is_forbidden_tensor(const std::string& name) { + static const std::string kTokenEmbd{"token_embd.weight"}; + if (name == kTokenEmbd) return true; + //if (auto pos = name.find("attn_kv_b.weight"); pos != std::string::npos) return true; + return false; +} +} + +bool iqk_should_modify_tensor([[maybe_unused]] const struct ggml_tensor * tensor) { + return false; + //if (is_forbidden_tensor(tensor->name)) return false; + //auto mptr = get_modify_info(tensor->type); + //return mptr ? true : false; +} + +bool iqk_modify_tensor(struct ggml_tensor * tensor) { + return false; + auto mptr = get_modify_info(tensor->type); + if (!mptr) return false; + if (is_forbidden_tensor(std::string{tensor->name})) return false; + + auto& m = *mptr; + int nrows = ggml_nrows(tensor); + int nchunks = nrows/m.nrows; + int max_thread = std::max(1, int(std::thread::hardware_concurrency()/2)); + int nthread = std::min(nchunks, max_thread); + auto row_size = ggml_row_size(tensor->type, tensor->ne[0]); + std::atomic counter(0); + auto compute = [&counter, &m, tensor, row_size, nchunks] () { + int64_t n_per_call = m.nrows*tensor->ne[0]; + while (true) { + int row = counter.fetch_add(1); + if (row >= nchunks) break; + m.mod_func(n_per_call, (char *)tensor->data + row_size*row*m.nrows); + } + }; + std::vector workers(nthread-1); + for (auto& w : workers) w = std::thread(compute); + compute(); + for (auto& w : workers) w.join(); + + return true; +} + +namespace { +const Repack * get_repack_info(ggml_type type) { + static const std::unordered_map k_map = { + { GGML_TYPE_IQ2_K, { GGML_TYPE_IQ2_K_R4, 4, (Repack::repack_func)repack_iq2_k} }, + { GGML_TYPE_IQ3_K, { GGML_TYPE_IQ3_K_R4, 4, (Repack::repack_func)repack_iq3_k} }, + { GGML_TYPE_IQ4_K, { GGML_TYPE_IQ4_K_R4, 4, (Repack::repack_func)repack_iq4_k} }, + { GGML_TYPE_IQ5_K, { GGML_TYPE_IQ5_K_R4, 4, (Repack::repack_func)repack_iq5_k} }, + { GGML_TYPE_IQ4_XS, { GGML_TYPE_IQ4_XS_R8, 8, (Repack::repack_func)repack_iq4_xs} }, + { GGML_TYPE_IQ4_KS, { GGML_TYPE_IQ4_KS_R4, 4, (Repack::repack_func)repack_iq4_ks} }, + { GGML_TYPE_IQ5_KS, { GGML_TYPE_IQ5_KS_R4, 4, (Repack::repack_func)repack_iq5_ks} }, + { GGML_TYPE_IQ4_NL, { GGML_TYPE_IQ4_NL_R4, 4, (Repack::repack_func)repack_iq4_nl} }, + { GGML_TYPE_IQ2_BN, { GGML_TYPE_IQ2_BN_R4, 4, (Repack::repack_func)repack_iq2_bn} }, + { GGML_TYPE_IQ2_XXS,{ GGML_TYPE_IQ2_XXS_R4,4, (Repack::repack_func)repack_iq2_xxs} }, + { GGML_TYPE_IQ2_XS, { GGML_TYPE_IQ2_XS_R4, 4, (Repack::repack_func)repack_iq2_xs} }, + { GGML_TYPE_IQ2_S, { GGML_TYPE_IQ2_S_R4, 4, (Repack::repack_func)repack_iq2_s} }, + { GGML_TYPE_IQ3_XXS,{ GGML_TYPE_IQ3_XXS_R4,4, (Repack::repack_func)repack_iq3_xxs} }, + { GGML_TYPE_IQ3_S, { GGML_TYPE_IQ3_S_R4, 4, (Repack::repack_func)repack_iq3_s} }, + { GGML_TYPE_Q2_K, { GGML_TYPE_Q2_K_R4, 4, (Repack::repack_func)repack_q2_k} }, + { GGML_TYPE_Q3_K, { GGML_TYPE_Q3_K_R4, 4, (Repack::repack_func)repack_q3_k} }, + { GGML_TYPE_Q4_K, { GGML_TYPE_Q4_K_R4, 4, (Repack::repack_func)repack_q4_k} }, + { GGML_TYPE_Q5_K, { GGML_TYPE_Q5_K_R4, 4, (Repack::repack_func)repack_q5_k} }, + { GGML_TYPE_Q6_K, { GGML_TYPE_Q6_K_R4, 4, (Repack::repack_func)repack_q6_k} }, + { GGML_TYPE_Q4_0, { GGML_TYPE_Q4_0_R8, 8, (Repack::repack_func)repack_q4_0} }, + { GGML_TYPE_Q5_0, { GGML_TYPE_Q5_0_R4, 4, (Repack::repack_func)repack_q5_0} }, + { GGML_TYPE_Q6_0, { GGML_TYPE_Q6_0_R4, 4, (Repack::repack_func)repack_q6_0} }, + { GGML_TYPE_Q8_0, { GGML_TYPE_Q8_0_R8, 8, (Repack::repack_func)repack_q8_0} }, + { GGML_TYPE_Q8_K, { GGML_TYPE_Q8_K_R8, 8, (Repack::repack_func)repack_q8_k} }, + { GGML_TYPE_Q8_KV, { GGML_TYPE_Q8_KV_R8, 8, (Repack::repack_func)repack_q8_KV} }, +#ifdef __AVX512BF16__ + { GGML_TYPE_BF16, { GGML_TYPE_BF16_R16, 16, (Repack::repack_func)repack_bf16}}, + { GGML_TYPE_F16, { GGML_TYPE_BF16_R16, 16, (Repack::repack_func)repack_bf16} }, +#endif + }; + auto it = k_map.find(type); + return it != k_map.end() ? &it->second : nullptr; +} +} + +int iqk_repacked_type(const struct ggml_tensor * tensor) { + if (!ggml_is_contiguous(tensor)) return (int)tensor->type; + if (is_forbidden_tensor(tensor->name)) return (int)tensor->type; + auto rptr = get_repack_info(tensor->type); + return rptr && tensor->ne[1] % rptr->num_rows == 0 ? (int)rptr->new_type : (int)tensor->type; +} + +void iqk_repack_tensor(struct ggml_tensor * tensor) { + constexpr int kChunk = 8; + if (!tensor) return; + if (!ggml_is_contiguous(tensor)) return; + if (is_forbidden_tensor(tensor->name)) return; + if (tensor->ne[1] % 4) return; + + auto rptr = get_repack_info(tensor->type); + if (!rptr) return; + if (tensor->ne[1] % rptr->num_rows) return; + + auto& r = *rptr; + + auto nrows = ggml_nrows(tensor); + + int max_thread = std::max(1, int(std::thread::hardware_concurrency()/2)); + int num_chunks = (nrows + kChunk*r.num_rows - 1)/(kChunk*r.num_rows); + int nthread = std::min(num_chunks, max_thread); + + //printf("%s(%s): %s -> %s. %d rows, %d chunks, %d threads\n", __func__, tensor->name, ggml_type_name(tensor->type), ggml_type_name(r.new_type), + // int(tensor->ne[1]), num_chunks, nthread); + + std::atomic counter(0);; + auto compute = [&counter, &r, tensor, num_chunks, chunkSize = kChunk] () { + int nrows = ggml_nrows(tensor); + int n_per_row = tensor->ne[0]; + auto row_size = ggml_row_size(tensor->type, n_per_row); + std::vector qtmp(r.num_rows*row_size); + auto data = (char *)tensor->data; + while (true) { + int chunk = counter.fetch_add(1); + if (chunk >= num_chunks) break; + int first_row = chunk*chunkSize*r.num_rows; + int last_row = std::min(first_row + chunkSize*r.num_rows, nrows); + for (int row = first_row; row < last_row; row += r.num_rows) { + std::memcpy(qtmp.data(), data + row*row_size, r.num_rows*row_size); + //r.repack(r.num_rows, n_per_row, qtmp.data(), data + row*row_size, true); + r.repack(r.num_rows, n_per_row, qtmp.data(), data + row*row_size, false); + } + } + }; + std::vector workers(nthread-1); + for (auto& w : workers) w = std::thread(compute); + compute(); + for (auto& w : workers) w.join(); + + tensor->type = r.new_type; +} + +void dequantize_row_ms_i2s(const void * vx, float * y, int64_t k) { + constexpr int kBlockSize = 128; + constexpr int kGroupSize = kBlockSize/4; + GGML_ASSERT(k % kBlockSize == 0); + const uint8_t * x = (const uint8_t *)vx; + const float * dptr = (const float *)(x + k/4); + const float d = dptr[0]; + int nb = k/kBlockSize; + for (int ib = 0; ib < nb; ++ib) { + for (int ig = 0; ig < kBlockSize/kGroupSize; ++ig) { + int shift = 6 - 2*ig; + for (int j = 0; j < kGroupSize; ++j) { + y[j] = d * (((x[j] >> shift) & 3) - 1); + } + y += kGroupSize; + } + x += kGroupSize; + } +} + +namespace { +template +class QuantizerIQKT { + static_assert(group_size == 8 || group_size == 4); + static_assert(block_size >= 8 && block_size%8 == 0); +public: + constexpr static int kSuperBlockSize = QK_K; + constexpr static int kBlockSize = block_size; + constexpr static int kGroupSize = group_size; + constexpr static int kNg = kBlockSize/kGroupSize; + constexpr static int kNblock = kSuperBlockSize/kBlockSize; + constexpr static int kNumVal = 1 << num_bits; // i.e, 16 bits per group of 8 + constexpr static float kScale = is_int ? 1.f : 31.75f; + constexpr static bool kVerbose = false; + + QuantizerIQKT(int num_clusters, int num_neighbours, int offset = 4096); + const float * values() const { return m_values.data(); } + + inline void find_best_match(float d, const float * xb, const float * weight, int * best_idx) const; + inline std::pair find_best_scale(const float * xb, const float * weight, const int * best_idx) const; + inline float find_best_inverse_scale(const float * xb, const float * weight, const int * best_idx) const; + + static inline void set_values(uint32_t i, float * result, float scale, int offset = 4096) { + uint32_t x = i + offset; + if constexpr (is_int) { + constexpr uint32_t ka = 0xCBAC1FED; + uint32_t s; + auto i8 = (const int8_t *)&s; + for (int k = 0; k < kGroupSize; ++k) { + x = ka*x; + s = x & 0x3f3f3f3f; + if constexpr (is_abs) { + result[k] = scale*std::abs(i8[0] + i8[1] + i8[2] + i8[3] - 126.f); + } else { + result[k] = scale*(i8[0] + i8[1] + i8[2] + i8[3] - 126.f); + } + } + } else { + constexpr uint32_t ka = 89226354; + constexpr uint32_t kb = 64248484; + constexpr uint32_t kmask = 0x8fff8fff; + constexpr uint32_t km32 = 0x3b603b60; + for (int k = 0; k < kGroupSize; ++k) { + x = ka*x + kb; + uint32_t s = (x & kmask) ^ km32; + float val = GGML_FP16_TO_FP32(s & 65535) + GGML_FP16_TO_FP32(s >> 16); + if constexpr (is_abs) result[k] = scale*std::abs(val); + else result[k] = scale*val; + } + } + } + + static inline int bin4(float x) { + if constexpr (is_abs) { + return x < 16.f ? 0 : x < 32.f ? 1 : x < 64.f ? 2 : 3; + } else { + return x < -24.f ? 0 : x < 0.0f ? 1 : x < 24.f ? 2 : 3; + } + } + static inline int bin5(float x) { + if constexpr (is_abs) { + return x < 11.2f ? 0 : x < 24.f ? 1 : x < 39.f ? 2 : x < 58.f ? 3 : 4; + } else { + return x < -48.f ? 0 : x < -16.f ? 1 : x < 16.f ? 2 : x < 48.f ? 3 : 4; + } + } + inline int bin3(int idim, float x) const { return x < m_mid[2*idim+0] ? 0 : x < m_mid[2*idim+1] ? 1 : 2; } + + static inline void set_weights(float sigma2_scale, int nblock, const float * x, const float * imatrix, float * row_weights) { + constexpr float kEps2 = 1e-14f; + constexpr float kWeight = 1e-4f; + for (int ibl = 0; ibl < nblock; ++ibl) { + + const float * xbl = x + ibl*kSuperBlockSize; + float * wbl = row_weights + ibl*kSuperBlockSize; + + float sumx2 = 0; + for (int j = 0; j < kSuperBlockSize; ++j) sumx2 += xbl[j]*xbl[j]; + if (sumx2 < kEps2*kSuperBlockSize) { + // all x in th super block are (almost) zero + for (int j = 0; j < kSuperBlockSize; ++j) wbl[j] = kWeight; + continue; + } + const float sigma2 = sigma2_scale*sumx2/kSuperBlockSize; + + if (imatrix) { + for (int ib = 0; ib < kSuperBlockSize/kBlockSize; ++ib) { + const float * qw = imatrix + ibl*kSuperBlockSize + ib*kBlockSize; + const float * xb = xbl + ib*kBlockSize; + float * wb = wbl + ib*kBlockSize; + float sumwx = 0, sumw2 = 0, sumx2 = 0; + for (int j = 0; j < kBlockSize; ++j) { + wb[j] = qw[j] * sqrtf(sigma2 + xb[j]*xb[j]); + sumwx += wb[j]*std::abs(xb[j]); + sumw2 += wb[j]*wb[j]; + sumx2 += xb[j]*xb[j]; + } + if (sumx2 < kEps2 || sumw2 < kEps2 || sumwx < kEps2) { + for (int j = 0; j < kBlockSize; ++j) wb[j] = kWeight; + } + } + } else { + for (int j = 0; j < kSuperBlockSize; ++j) wbl[j] = 0.25f*sigma2 + xbl[j]*xbl[j]; + } + } + } +private: + static std::vector cluster_points(const std::vector& points, int ncluster, int niter, float * mid); + static std::vector> finalize_clusters(int num_neighbours, const std::vector& points, const std::vector& clusters, + std::vector>& c_values); + std::vector m_values; + std::vector m_clusters; + std::vector> m_in_cluster; + std::vector> m_c_values; + float m_mid[4*kGroupSize]; +}; + +template +QuantizerIQKT::QuantizerIQKT(int num_clusters, int num_neighbours, int offset) { + m_values.resize(kNumVal*kGroupSize); + float * data = m_values.data(); + for (int i = 0; i < kNumVal; ++i) { + set_values(i, data, kScale, offset); + data += kGroupSize; + } + if (num_clusters == 0) return; + // Make 128 clusters. + // Note: we get a slightly better result by using 64 clusters + // at the expense of almost doubling the quantization time. + m_clusters = cluster_points(m_values, num_clusters, 200, m_mid); + GGML_ASSERT(!m_clusters.empty()); + m_in_cluster = finalize_clusters(num_neighbours, m_values, m_clusters, m_c_values); +} + +template +std::pair QuantizerIQKT::find_best_scale( + const float * xb, const float * weight, const int * best_idx) const { + float sumqx = 0, sumq2 = 0; +#ifdef __AVX2__ + auto vqx = _mm256_setzero_ps(); + auto vq2 = _mm256_setzero_ps(); + for (int l = 0; l < kBlockSize; l += 8) { + auto vx = _mm256_loadu_ps(xb+l); + auto vw = _mm256_loadu_ps(weight+l); + auto vq = kGroupSize == 8 ? _mm256_loadu_ps(m_values.data() + kGroupSize*best_idx[l/kGroupSize]) : + _mm256_set_m128(_mm_loadu_ps(m_values.data() + kGroupSize*best_idx[l/kGroupSize+1]), + _mm_loadu_ps(m_values.data() + kGroupSize*best_idx[l/kGroupSize+0])); + auto vqw = _mm256_mul_ps(vq, vw); + vqx = _mm256_fmadd_ps(vqw, vx, vqx); + vq2 = _mm256_fmadd_ps(vqw, vq, vq2); + } + sumqx = hsum_float_8(vqx); + sumq2 = hsum_float_8(vq2); +#else + for (int l = 0; l < kNg; ++l) { + auto xl = xb + kGroupSize*l; + auto wl = weight + kGroupSize*l; + auto ql = m_values.data() + kGroupSize*best_idx[l]; + for (int k = 0; k < kGroupSize; ++k) { + sumqx += wl[k]*ql[k]*xl[k]; + sumq2 += wl[k]*ql[k]*ql[k]; + } + } +#endif + return sumq2 > 0 ? std::make_pair(sumqx/sumq2, sumqx*sumqx/sumq2) : std::make_pair(0.f, 0.f); +} + +template +float QuantizerIQKT::find_best_inverse_scale( + const float * xb, const float * weight, const int * best_idx) const { + float sumqx = 0, sumx2 = 0; +#ifdef __AVX2__ + auto vqx = _mm256_setzero_ps(); + auto vx2 = _mm256_setzero_ps(); + for (int l = 0; l < kBlockSize; l += 8) { + auto vx = _mm256_loadu_ps(xb+l); + auto vw = _mm256_loadu_ps(weight+l); + auto vq = kGroupSize == 8 ? _mm256_loadu_ps(m_values.data() + kGroupSize*best_idx[l/kGroupSize]) : + _mm256_set_m128(_mm_loadu_ps(m_values.data() + kGroupSize*best_idx[l/kGroupSize+1]), + _mm_loadu_ps(m_values.data() + kGroupSize*best_idx[l/kGroupSize+0])); + auto vxw = _mm256_mul_ps(vx, vw); + vx2 = _mm256_fmadd_ps(vxw, vx, vx2); + vqx = _mm256_fmadd_ps(vxw, vq, vqx); + } + sumqx = hsum_float_8(vqx); + sumx2 = hsum_float_8(vx2); +#else + for (int l = 0; l < kNg; ++l) { + auto xl = xb + kGroupSize*l; + auto wl = weight + kGroupSize*l; + auto ql = m_values.data() + kGroupSize*best_idx[l]; + for (int k = 0; k < kGroupSize; ++k) { + sumqx += wl[k]*ql[k]*xl[k]; + sumx2 += wl[k]*xl[k]*xl[k]; + } + } +#endif + return sumx2 > 0 ? sumqx/sumx2 : 0.f; +} + +template +void QuantizerIQKT::find_best_match(float d, + [[maybe_unused]] const float * xb, [[maybe_unused]] const float * weight, int * best_idx) const { + if (!d) { + std::memset(best_idx, 0, kNg*sizeof(int)); + return; + } + [[maybe_unused]] int ncluster = m_clusters.size()/kGroupSize; + [[maybe_unused]] float id = 1/d; +#ifdef __AVX2__ + if constexpr (kGroupSize == 8) { + __m256 sqx[8]; + const __m256i add_idx = _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0); + float sx[8]; + int index[8]; + auto vid = _mm256_set1_ps(id); + auto add8 = _mm256_set1_epi32(8); + for (int l = 0; l < kNg; ++l) { + auto xl = xb + 8*l; + auto wl = weight + 8*l; + auto vx = _mm256_mul_ps(vid, _mm256_loadu_ps(xl)); + auto vw = _mm256_loadu_ps(wl); + int jbest = -1; + if (kGroupSize == 8 && (ncluster == 256 || ncluster == 6561)) { + _mm256_store_ps(sx, vx); + uint16_t u = 0; + if (ncluster == 256) { + for (int j = 0; j < 8; ++j) if (sx[j] > m_mid[j]) u |= (1 << j); + } else { + int s = 1; + for (int j = 0; j < 8; ++j) { u += s*bin3(j, sx[j]); s *= 3; } + } + jbest = u; + } else { + auto vbest = _mm256_set1_ps(INFINITY); + auto best_index = _mm256_set1_epi32(-1); + float best = INFINITY; + auto idx = add_idx; + for (int j = 0; j < ncluster; j += 8) { + for (int i = 0; i < 8; ++i) { + auto vq = _mm256_loadu_ps(m_clusters.data() + kGroupSize*(j+i)); + auto vdiff = _mm256_sub_ps(vq, vx); + sqx[i] = _mm256_mul_ps(vw, _mm256_mul_ps(vdiff, vdiff)); + } + auto score = hsum_float_8x8(sqx); + auto mask = _mm256_cmp_ps(score, vbest, _CMP_LT_OQ); + best_index = _mm256_or_si256(_mm256_and_si256(_mm256_castps_si256(mask), idx), + _mm256_andnot_si256(_mm256_castps_si256(mask), best_index)); + vbest = _mm256_min_ps(vbest, score); + idx = _mm256_add_epi32(idx, add8); + } + _mm256_store_ps(sx, vbest); + _mm256_store_si256((__m256i *)index, best_index); + for (int i = 0; i < 8; ++i) { + if (sx[i] < best) { best = sx[i]; jbest = index[i]; } + } + } + auto& points = m_in_cluster[jbest]; + auto& values = points.empty() ? m_values : m_c_values[jbest]; + int npoint = values.size()/kGroupSize; + GGML_ASSERT(npoint > 0 && npoint%8 == 0); + int jbest_cluster = jbest; + auto vbest = _mm256_set1_ps(INFINITY); + auto best_index = _mm256_set1_epi32(-1); + auto best = INFINITY; jbest = -1; + auto idx = add_idx; + for (int j = 0; j < npoint; j += 8) { + for (int i = 0; i < 8; ++i) { + auto vq = _mm256_loadu_ps(values.data() + kGroupSize*(j+i)); + auto vdiff = _mm256_sub_ps(vq, vx); + sqx[i] = _mm256_mul_ps(vw, _mm256_mul_ps(vdiff, vdiff)); + } + auto score = hsum_float_8x8(sqx); + auto mask = _mm256_cmp_ps(score, vbest, _CMP_LT_OQ); + best_index = _mm256_or_si256(_mm256_and_si256(_mm256_castps_si256(mask), idx), + _mm256_andnot_si256(_mm256_castps_si256(mask), best_index)); + vbest = _mm256_min_ps(vbest, score); + idx = _mm256_add_epi32(idx, add8); + } + _mm256_store_ps(sx, vbest); + _mm256_store_si256((__m256i *)index, best_index); + for (int i = 0; i < 8; ++i) { + if (sx[i] < best) { best = sx[i]; jbest = index[i]; } + } + if (jbest < 0) { + fprintf(stderr, "Oops: jbest = %d for cluster %d with %d points\n", jbest, jbest_cluster, int(points.size())); + GGML_ASSERT(false); + } + best_idx[l] = points.empty() ? jbest : points[jbest]; + } + } else { + __m256 sqx[4]; + const __m256i add_idx = _mm256_set_epi32(7, 5, 3, 1, 6, 4, 2, 0); + const __m256 sign_bit = _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffff)); + float sx[8]; + int index[8]; + auto vid_p = _mm256_set1_ps(id); + auto add8 = _mm256_set1_epi32(8); + for (int l = 0; l < kNg; ++l) { + auto xl = xb + 4*l; + auto wl = weight + 4*l; + auto vx4 = _mm_loadu_ps(xl); + auto vx = _mm256_mul_ps(vid_p, _mm256_set_m128(vx4, vx4)); + auto vw4 = _mm_loadu_ps(wl); + auto vw = _mm256_set_m128(vw4, vw4); + int jbest = -1; + if (ncluster == 256 || ncluster == 625) { + _mm256_storeu_ps(sx, vx); + uint16_t u = 0; + if (ncluster == 256) { + for (int k = 0; k < 4; ++k) u |= (bin4(sx[k]) << 2*k); + } else { + int l = 1; + for (int k = 0; k < 4; ++k) { u += bin5(sx[k])*l; l *= 5; } + } + jbest = u; + } else { + auto vbest = _mm256_set1_ps(INFINITY); + auto best_index = _mm256_set1_epi32(-1); + float best = INFINITY; + auto idx = add_idx; + for (int j = 0; j < ncluster; j += 8) { + for (int i = 0; i < 4; ++i) { + auto vq = _mm256_loadu_ps(m_clusters.data() + kGroupSize*(j+2*i)); + auto vdiff = _mm256_sub_ps(vq, vx); + vdiff = _mm256_and_ps(sign_bit, vdiff); + sqx[i] = _mm256_mul_ps(vw, _mm256_mul_ps(vdiff, _mm256_mul_ps(vdiff, vdiff))); + } + auto score = hsum_float_4x8(sqx); + auto mask = _mm256_cmp_ps(score, vbest, _CMP_LT_OQ); + best_index = _mm256_or_si256(_mm256_and_si256(_mm256_castps_si256(mask), idx), + _mm256_andnot_si256(_mm256_castps_si256(mask), best_index)); + vbest = _mm256_min_ps(vbest, score); + idx = _mm256_add_epi32(idx, add8); + } + _mm256_store_ps(sx, vbest); + _mm256_store_si256((__m256i *)index, best_index); + for (int i = 0; i < 8; ++i) { + if (sx[i] < best) { best = sx[i]; jbest = index[i]; } + } + } + auto& points = m_in_cluster[jbest]; + auto& values = m_c_values[jbest]; + GGML_ASSERT(!points.empty() && points.size()%8 == 0); + int jbest_cluster = jbest; + auto vbest = _mm256_set1_ps(INFINITY); + auto best_index = _mm256_set1_epi32(-1); + float best = INFINITY; jbest = -1; + auto idx = add_idx; + for (int j = 0; j < int(points.size()); j += 8) { + for (int i = 0; i < 4; ++i) { + auto vq = _mm256_loadu_ps(values.data() + kGroupSize*(j+2*i)); + auto vdiff = _mm256_sub_ps(vq, vx); + sqx[i] = _mm256_mul_ps(vw, _mm256_mul_ps(vdiff, vdiff)); + } + auto score = hsum_float_4x8(sqx); + auto mask = _mm256_cmp_ps(score, vbest, _CMP_LT_OQ); + best_index = _mm256_or_si256(_mm256_and_si256(_mm256_castps_si256(mask), idx), + _mm256_andnot_si256(_mm256_castps_si256(mask), best_index)); + vbest = _mm256_min_ps(vbest, score); + idx = _mm256_add_epi32(idx, add8); + } + _mm256_store_ps(sx, vbest); + _mm256_store_si256((__m256i *)index, best_index); + for (int i = 0; i < 8; ++i) { + if (sx[i] < best) { best = sx[i]; jbest = index[i]; } + } + if (jbest < 0) { + fprintf(stderr, "Oops: jbest = %d for cluster %d with %d points\n", jbest, jbest_cluster, int(points.size())); + GGML_ASSERT(false); + } + best_idx[l] = points[jbest]; + } + } +#else + // TODO + std::memset(best_idx, 0, kNg*sizeof(int)); +#endif +} + +template +std::vector> QuantizerIQKT::finalize_clusters(int num_neighbours, + const std::vector& values, const std::vector& clusters, std::vector>& c_values) { + int ncluster = clusters.size()/kGroupSize; + std::vector> p_in_cluster(ncluster); + std::vector which_cluster(num_neighbours*kNumVal); + std::vector ibest(num_neighbours); + std::vector best(num_neighbours); + for (int ip = 0; ip < kNumVal; ++ip) { + auto vp = values.data() + ip*kGroupSize; + for (int j = 0; j < num_neighbours; ++j) { + best[j] = INFINITY; ibest[j] = -1; + } + for (int ic = 0; ic < ncluster; ++ic) { + auto vc = clusters.data() + ic*kGroupSize; + float dist2 = 0; + for (int k = 0; k < kGroupSize; ++k) { + float d = vp[k] - vc[k]; dist2 += d*d; + } + for (int j = 0; j < num_neighbours; ++j) { + if (dist2 < best[j]) { + for (int k = num_neighbours-1; k > j; --k) { + best[k] = best[k-1]; ibest[k] = ibest[k-1]; + } + best[j] = dist2; ibest[j] = ic; + break; + } + } + } + for (int j = 0; j < num_neighbours; ++j) { + if (ibest[j] < 0) { + printf("Oops: ibest[%d] = %d\n", j, ibest[j]); + } + GGML_ASSERT(ibest[j] >= 0); + p_in_cluster[ibest[j]].push_back(ip); + } + std::memcpy(which_cluster.data() + num_neighbours*ip, ibest.data(), num_neighbours*sizeof(int)); + } + std::vector> extra; + extra.reserve(kNumVal); + for (int ic = 0; ic < ncluster; ++ic) { + auto& points = p_in_cluster[ic]; + if (!points.empty() && points.size()%8 == 0) continue; + extra.clear(); + auto vc = clusters.data() + ic*kGroupSize; + for (int ip = 0; ip < kNumVal; ++ip) { + bool can_add = true; + for (int j = 0; j < num_neighbours; ++j) { + if (which_cluster[num_neighbours*ip+j] == ic) { can_add = false; break; } + } + if (!can_add) continue; + auto vp = values.data() + ip*kGroupSize; + float dist2 = 0; + for (int k = 0; k < kGroupSize; ++k) { + float d = vp[k] - vc[k]; dist2 += d*d; + } + extra.push_back(std::make_pair(dist2, ip)); + } + std::sort(extra.begin(), extra.end()); + int nadd = 8*((points.size()+7)/8) - points.size(); + for (int i = 0; i < nadd; ++i) points.push_back(extra[i].second); + GGML_ASSERT(points.size()%8 == 0); + } + auto min = p_in_cluster.front().size(), max = p_in_cluster.front().size(); + for (auto& points : p_in_cluster) { + min = std::min(min, points.size()); + max = std::max(max, points.size()); + } + c_values.resize(p_in_cluster.size()); + for (int i = 0; i < int(p_in_cluster.size()); ++i) { + auto& points = p_in_cluster[i]; + c_values[i].resize(points.size()*kGroupSize); + auto ptr = c_values[i].data(); + for (auto j : points) { + std::memcpy(ptr, values.data() + j*kGroupSize, kGroupSize*sizeof(float)); + ptr += kGroupSize; + } + } + + if (kVerbose) { + printf("%s: prepared %d clusters\n", __func__, ncluster); + printf(" min number of points in a cluster: %d\n", int(min)); + printf(" max number of points in a cluster: %d\n", int(max)); + } + return p_in_cluster; +} + +template +std::vector QuantizerIQKT::cluster_points(const std::vector& points, int ncluster, int niter, float * mid) { + constexpr int ndim = kGroupSize; + GGML_ASSERT(points.size() % ndim == 0); + int npoint = points.size() / ndim; + GGML_ASSERT(npoint >= 2*ncluster); + std::vector> range(ndim, std::make_pair(INFINITY, -INFINITY)); + double Fo = 0; + for (int i = 0; i < npoint; ++i) { + auto v = points.data() + i*ndim; + for (int k = 0; k < ndim; ++k) { + Fo += v[k]*v[k]; + range[k].first = std::min(range[k].first, v[k]); + range[k].second = std::max(range[k].second, v[k]); + } + } + if (kVerbose) printf("%s (ndim = %d, npoint = %d): Fo = %g\n", __func__, ndim, npoint, Fo/points.size()); + if constexpr (is_abs) { + std::vector P(npoint); + for (int idim = 0; idim < ndim; ++idim) { + for (int ip = 0; ip < npoint; ++ip) P[ip] = points[ip*ndim+idim]; + std::sort(P.begin(), P.end()); + if (ndim == 8 && ncluster == 6561) { + mid[2*idim + 0] = P[npoint/3]; + mid[2*idim + 1] = P[2*npoint/3]; + } else { + mid[idim] = npoint%2 == 0 ? 0.5f*(P[npoint/2] + P[npoint/2-1]) : P[npoint/2]; + if (kVerbose) printf("%s: mid[%d] = %g\n", __func__, idim, mid[idim]); + } + } + } else { + for (int k = 0; k < ndim; ++k) mid[k] = 0.5f*(range[k].first + range[k].second); + } + std::vector sump(ncluster*ndim); + std::vector counts(ncluster); + std::vector result(ncluster*ndim); + if (ndim == 8 && (ncluster == 256 || ncluster == 6561)) { + std::memset(sump.data(), 0, sump.size()*sizeof(float)); + std::memset(counts.data(), 0, counts.size()*sizeof(int)); + for (int ip = 0; ip < npoint; ++ip) { + auto vp = points.data() + ndim*ip; + uint16_t u = 0; + if (ncluster == 256) { + for (int k = 0; k < ndim; ++k) if (vp[k] > mid[k]) u |= (1 << k); + } else { + int s = 1; + for (int k = 0; k < ndim; ++k) { + int bin = vp[k] < mid[2*k+0] ? 0 : vp[k] < mid[2*k+1] ? 1 : 2; + u += s*bin; s *= 3; + } + } + ++counts[u]; + for (int k = 0; k < ndim; ++k) sump[ndim*u + k] += vp[k]; + } + for (int ic = 0; ic < ncluster; ++ic) { + if (!counts[ic]) { + printf("%s: Oops. Cluster %d has no points\n", __func__, ic); + GGML_ABORT("fatal error"); + } + for (int k = 0; k < ndim; ++k) result[ic*ndim + k] = sump[ic*ndim + k]/counts[ic]; + } + return result; + } + else if (ndim == 4 && (ncluster == 256 || ncluster == 625)) { + std::memset(sump.data(), 0, sump.size()*sizeof(float)); + std::memset(counts.data(), 0, counts.size()*sizeof(int)); + for (int ip = 0; ip < npoint; ++ip) { + auto vp = points.data() + ndim*ip; + uint16_t u = 0; + if (ncluster == 256) { + for (int k = 0; k < ndim; ++k) u |= (bin4(vp[k]) << 2*k); + } else { + int s = 1; + for (int k = 0; k < ndim; ++k) { u += s*bin5(vp[k]); s *= 5; } + } + if (u >= int(counts.size())) { + printf("Oops: u = %u, vp = %g, %g, %g, %g\n", u, vp[0], vp[1], vp[2], vp[3]); + u = 0; + if (ncluster == 256) { + for (int k = 0; k < ndim; ++k) { + auto bin = bin4(vp[k]); u |= (bin << 2*k); + printf(" bin[%d] = %d, u = %u", k, bin, u); + } + } else { + for (int k = 0; k < ndim; ++k) printf(" bin[%d] = %d", k, bin5(vp[k])); + } + printf("\n"); + GGML_ABORT("fatal error"); + } + ++counts[u]; + for (int k = 0; k < ndim; ++k) sump[ndim*u + k] += vp[k]; + } + int nzero = 0; + for (int ic = 0; ic < ncluster; ++ic) { + if (!counts[ic]) { + ++nzero; + printf("%s: Oops. Cluster %d has no points: ", __func__, ic); + for (int k = 0; k < ndim; ++k) { + int l = (ic >> 2*k) & 3; + printf(" %d", l); + } + printf("\n"); + } else { + for (int k = 0; k < ndim; ++k) result[ic*ndim + k] = sump[ic*ndim + k]/counts[ic]; + } + } + if (nzero > 0) printf("%s: %d out of %d clusters dir not have any points\n", __func__, nzero, ncluster); + return result; + } + std::mt19937 rndm(1234); + float scale = 1.f/4294967296.f; + for (int i = 0; i < ncluster; ++i) { + auto v = result.data() + i*ndim; + for (int k = 0; k < ndim; ++k) v[k] = range[k].first + (range[k].second - range[k].first)*scale*rndm(); + } + std::vector which_cluster(npoint, -1); + double Flast = Fo; + for (int iter = 0; iter < niter; ++iter) { + std::memset(sump.data(), 0, sump.size()*sizeof(float)); + std::memset(counts.data(), 0, counts.size()*sizeof(int)); + int nchanged = 0; + double F = 0; + for (int ip = 0; ip < npoint; ++ip) { + auto vp = points.data() + ndim*ip; + float best = INFINITY; int ibest = -1; + for (int ic = 0; ic < ncluster; ++ic) { + auto vc = result.data() + ndim*ic; + float dist2 = 0; + for (int k = 0; k < ndim; ++k) { + float d = vp[k] - vc[k]; dist2 += d*d; + } + if (dist2 < best) { + best = dist2; ibest = ic; + } + } + if (ibest < 0) { + printf("Oops(iteration %d) - failed to find cluster for point", iter); + for (int k = 0; k < ndim; ++k) printf(" %g", vp[k]); + printf("\nHave %d clusters\n", ncluster); + } + GGML_ASSERT(ibest >= 0); + F += best; + if (which_cluster[ip] != ibest) ++nchanged; + which_cluster[ip] = ibest; + ++counts[ibest]; + auto vc = sump.data() + ndim*ibest; + for (int k = 0; k < ndim; ++k) vc[k] += vp[k]; + } + if (nchanged == 0) break; + for (int ic = 0; ic < ncluster; ++ic) { + float norm = counts[ic] > 0 ? 1.f/counts[ic] : 0.f; + auto vc = sump.data() + ndim*ic; + auto r = result.data() + ndim*ic; + for (int k = 0; k < ndim; ++k) r[k] = vc[k]*norm; + } + if (kVerbose) printf("%s(iteration %d): F = %g, nchanged = %d\n", __func__, iter+1, F/points.size(), nchanged); + if (iter > 1 && Flast/F - 1 < 1e-6) break; + Flast = F; + } + int nzero = 0; + for (int ic = 0; ic < ncluster; ++ic) { + if (!counts[ic]) ++nzero; + } + if (nzero > 0) printf("%s: there are %d empty clusters\n", __func__, nzero); + return result; +} + +// ========================================== iq1_kt ==================================================== + +using QuantizerIQ1KT = QuantizerIQKT<32, 8, 13, false, true>; + +const QuantizerIQ1KT& iq1kt_quantizer() { + static std::mutex mutex; + static std::unique_ptr quantizer; + std::lock_guard lock(mutex); + if (!quantizer) quantizer = std::make_unique(256, 32); + return *quantizer; +} + +void quantize_row_iq1_kt_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, float * all_scales, float * all_weights, + int * all_idx) { + + constexpr float kSigmaScale = 2.0f; + using Q = QuantizerIQ1KT; + + static_assert(Q::kNumVal%8 == 0); + + float * dptr = (float *)vy; + + block_iq1_kt * y = (block_iq1_kt *)(dptr + 1); + + int best_idx[2*Q::kNg]; + + auto& quantizer = iq1kt_quantizer(); + + int nblock = n_per_row / Q::kSuperBlockSize; + + Q::set_weights(kSigmaScale, nblock, x, quant_weights, all_weights); + + float amax_row = 0; + for (int j = 0; j < n_per_row; ++j) { + amax_row = std::max(amax_row, std::abs(x[j])); + } + + float amax_scale = 0, max_scale = 0; + + for (int ibl = 0; ibl < nblock; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq1_kt)); + + const float * xbl = x + ibl*Q::kSuperBlockSize; + auto scales = all_scales + ibl*Q::kNblock; + + for (int ib = 0; ib < Q::kNblock; ++ib) { + const float * xb = xbl + Q::kBlockSize*ib; + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + float amax = 0; + for (int j = 0; j < Q::kBlockSize; ++j) { + float ax = std::abs(xb[j]); + amax = std::max(amax, ax); + } + if (amax < 1e-16f) { + scales[ib] = 0.0f; + for (int ig = 0; ig < Q::kNg; ++ig) all_idx[(ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize + ig] = 0; + continue; + } + float scale_0 = std::max(90.f, 124.f*amax/amax_row); + quantizer.find_best_match( amax/scale_0, xb, weight, best_idx); + auto [dp, score_p] = quantizer.find_best_scale(xb, weight, best_idx); + quantizer.find_best_match(-amax/scale_0, xb, weight, best_idx + Q::kNg); + auto [dm, score_m] = quantizer.find_best_scale(xb, weight, best_idx + Q::kNg); + + auto idx = best_idx; + if (score_p > score_m) scales[ib] = dp; + else { + scales[ib] = dm; idx += Q::kNg; score_p = score_m; + } + for (int ig = 0; ig < Q::kNg; ++ig) all_idx[(ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize + ig] = idx[ig]; + + scale_0 -= 8; + quantizer.find_best_match( amax/scale_0, xb, weight, best_idx); + auto [dp1, score_p1] = quantizer.find_best_scale(xb, weight, best_idx); + quantizer.find_best_match(-amax/scale_0, xb, weight, best_idx + Q::kNg); + auto [dm1, score_m1] = quantizer.find_best_scale(xb, weight, best_idx + Q::kNg); + + if (score_p1 > score_p || score_m1 > score_p) { + idx = best_idx; + if (score_p1 > score_m1) scales[ib] = dp1; + else { + scales[ib] = dm1; idx += Q::kNg; + } + for (int ig = 0; ig < Q::kNg; ++ig) all_idx[(ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize + ig] = idx[ig]; + } + + float abs_scale = std::abs(scales[ib]); + if (abs_scale > amax_scale) { + amax_scale = abs_scale; + max_scale = scales[ib]; + } + } + + } + + if (!max_scale) { + *dptr = 0; + return; + } + + float d = max_scale/iq4k_values[0]; + float best = 0; + for (int itry = -9; itry <= 9; ++itry) { + float id = (itry + iq4k_values[0])/max_scale; + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * xb = x + ibl*Q::kSuperBlockSize; + const float * wb = all_weights + ibl*Q::kSuperBlockSize; + auto scales = all_scales + ibl*Q::kNblock; + for (int ib = 0; ib < Q::kNblock; ++ib) { + int ls = best_index_iq4nl(iq4k_values, id*scales[ib]); + float dl = iq4k_values[ls]; + for (int ig = 0; ig < Q::kNg; ++ig) { + auto qb = quantizer.values() + Q::kGroupSize*all_idx[(ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize + ig]; + for (int j = 0; j < Q::kGroupSize; ++j) { + int jj = ig*Q::kGroupSize + j; + float q = dl*qb[j]; + sumqx += wb[jj]*xb[jj]*q; + sumq2 += wb[jj]*q*q; + } + } + xb += Q::kBlockSize; + wb += Q::kBlockSize; + } + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; + } + } + + float id = d ? 1/d : 0.f; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto scales = all_scales + ibl*Q::kNblock; + for (int ib = 0; ib < Q::kNblock; ++ib) { + int ls = best_index_iq4nl(iq4k_values, id*scales[ib]); + y[ibl].sh[ib] = ls; + } + } + + *dptr = d; + if (!d) return; + + for (int iloop = 0; iloop < 1; ++iloop) { + + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + + const float * xbl = x + ibl*Q::kSuperBlockSize; + + for (int ib = 0; ib < Q::kNblock; ++ib) { + const float * xb = xbl + Q::kBlockSize*ib; + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + int ls = iq4k_values[y[ibl].sh[ib] & 0xf]; + float dl = d*ls; + quantizer.find_best_match(dl, xb, weight, best_idx); + + auto prev_idx = all_idx + (ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize; + + float mse1 = 0, mse2 = 0; + for (int ig = 0; ig < Q::kNg; ++ig) { + auto q1 = quantizer.values() + Q::kGroupSize*prev_idx[ig]; + auto q2 = quantizer.values() + Q::kGroupSize*best_idx[ig]; + for (int j = 0; j < Q::kGroupSize; ++j) { + int jj = ig*Q::kGroupSize + j; + float diff1 = xb[jj] - dl*q1[j]; + float diff2 = xb[jj] - dl*q2[j]; + mse1 += weight[jj]*diff1*diff1; + mse2 += weight[jj]*diff2*diff2; + } + } + if (mse1 < mse2) { + for (int ig = 0; ig < Q::kNg; ++ig) best_idx[ig] = prev_idx[ig]; + } else { + for (int ig = 0; ig < Q::kNg; ++ig) prev_idx[ig] = best_idx[ig]; + } + + for (int j = 0; j < Q::kNg; ++j) { + y[ibl].ql[ib*Q::kNg+j] = best_idx[j] & 0xff; + y[ibl].qh[(ib%(Q::kNblock/2))*Q::kNg+j] |= (((best_idx[j] >> 8) & 0xf) << 4*(ib/(Q::kNblock/2))); + y[ibl].sh[ib] |= ((best_idx[j] >> 12) << (4+j)); + auto xl = xb + Q::kGroupSize*j; + auto wl = weight + Q::kGroupSize*j; + auto ql = quantizer.values() + best_idx[j]*Q::kGroupSize; + for (int k = 0; k < Q::kGroupSize; ++k) { + float q = ql[k]*ls; + sumqx += wl[k]*xl[k]*q; + sumq2 += wl[k]*q*q; + } + } + } + } + if (sumq2 > 0) { + d = sumqx/sumq2; + *dptr = d * 1.07f; + if (!d) return; + } else { + break; + } + + } + +} +} + +void quantize_row_iq1_kt_ref(const float * GGML_RESTRICT x, block_iq1_kt * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq1_kt(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq1_kt(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_iq1_kt * y = (block_iq1_kt *)vy; + quantize_row_iq1_kt_ref(x, y, k); +} + +size_t quantize_iq1_kt(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ1_KT, n_per_row); + std::vector scales(n_per_row/QuantizerIQ1KT::kBlockSize); + std::vector weights(n_per_row); + std::vector idx(n_per_row/QuantizerIQ1KT::kGroupSize); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrows; ++row) { + quantize_row_iq1_kt_impl(src, (void *)qrow, n_per_row, imatrix, scales.data(), weights.data(), idx.data()); + src += n_per_row; + qrow += row_size; + } + return nrows * row_size; +} + +void dequantize_row_iq1_kt(const block_iq1_kt * x, float * y, int64_t k) { + assert(k % QuantizerIQ1KT::kSuperBlockSize == 0); + using Q = QuantizerIQ1KT; + const int nb = k / Q::kSuperBlockSize; + const float * dptr = (const float *)x; + const float d = *dptr * Q::kScale; + x = (const block_iq1_kt *)(dptr + 1); + auto& deq = iq1kt_quantizer(); + for (int ibl = 0; ibl < nb; ++ibl) { + for (int ib = 0; ib < Q::kNblock; ++ib) { + float sl = d * iq4k_values[x[ibl].sh[ib] & 0xf]; + for (int ig = 0; ig < Q::kNg; ++ig) { + uint16_t idx = x[ibl].ql[ib*Q::kNg + ig] | ((x[ibl].qh[(ib%(Q::kNblock/2))*Q::kNg + ig] << (8 - 4*(ib/(Q::kNblock/2)))) & 0xf00); + idx |= (x[ibl].sh[ib] << (8 - ig) & 0x1000); + deq.set_values(idx, y, sl); + y += Q::kGroupSize; + } + } + } +} + +void vec_dot_iq1_kt_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ1_KT, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + +} + +// ========================================== iq2_kt ==================================================== + +namespace { + +using QuantizerIQ2KT = QuantizerIQKT<32, 8, 16, false, true>; + +const QuantizerIQ2KT& iq2kt_quantizer() { + static std::mutex mutex; + static std::unique_ptr quantizer; + std::lock_guard lock(mutex); + if (!quantizer) quantizer = std::make_unique(256, 8); + return *quantizer; +} + +void quantize_row_iq2_kt_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, float * all_scales, float * all_weights, + int * all_idx) { + + constexpr float kSigmaScale = 2.0f; + using Q = QuantizerIQ2KT; + + static_assert(Q::kNumVal%8 == 0); + + float * dptr = (float *)vy; + + block_iq2_kt * y = (block_iq2_kt *)(dptr + 1); + + int best_idx[2*Q::kNg]; + + auto& quantizer = iq2kt_quantizer(); + + int nblock = n_per_row / Q::kSuperBlockSize; + + Q::set_weights(kSigmaScale, nblock, x, quant_weights, all_weights); + + float amax_row = 0; + for (int j = 0; j < n_per_row; ++j) { + amax_row = std::max(amax_row, std::abs(x[j])); + } + + float amax_scale = 0, max_scale = 0; + + for (int ibl = 0; ibl < nblock; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq2_kt)); + + const float * xbl = x + ibl*Q::kSuperBlockSize; + auto scales = all_scales + ibl*Q::kNblock; + + for (int ib = 0; ib < Q::kNblock; ++ib) { + const float * xb = xbl + Q::kBlockSize*ib; + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + float amax = 0; + for (int j = 0; j < Q::kBlockSize; ++j) { + float ax = std::abs(xb[j]); + amax = std::max(amax, ax); + } + if (amax < 1e-16f) { + scales[ib] = 0.0f; + for (int ig = 0; ig < Q::kNg; ++ig) all_idx[(ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize + ig] = 0; + continue; + } + float scale_0 = std::max(90.f, 124.f*amax/amax_row); + quantizer.find_best_match( amax/scale_0, xb, weight, best_idx); + auto [dp, score_p] = quantizer.find_best_scale(xb, weight, best_idx); + quantizer.find_best_match(-amax/scale_0, xb, weight, best_idx + Q::kNg); + auto [dm, score_m] = quantizer.find_best_scale(xb, weight, best_idx + Q::kNg); + + auto idx = best_idx; + if (score_p > score_m) scales[ib] = dp; + else { + scales[ib] = dm; idx += Q::kNg; + } + for (int ig = 0; ig < Q::kNg; ++ig) all_idx[(ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize + ig] = idx[ig]; + + float abs_scale = std::abs(scales[ib]); + if (abs_scale > amax_scale) { + amax_scale = abs_scale; + max_scale = scales[ib]; + } + } + + } + + if (!max_scale) { + *dptr = 0; + return; + } + + float d = max_scale/iq4k_values[0]; + float best = 0; + for (int itry = -9; itry <= 9; ++itry) { + float id = (itry + iq4k_values[0])/max_scale; + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * xb = x + ibl*Q::kSuperBlockSize; + const float * wb = all_weights + ibl*Q::kSuperBlockSize; + auto scales = all_scales + ibl*Q::kNblock; + for (int ib = 0; ib < Q::kNblock; ++ib) { + int ls = best_index_iq4nl(iq4k_values, id*scales[ib]); + float dl = iq4k_values[ls]; + for (int ig = 0; ig < Q::kNg; ++ig) { + auto qb = quantizer.values() + Q::kGroupSize*all_idx[(ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize + ig]; + for (int j = 0; j < Q::kGroupSize; ++j) { + int jj = ig*Q::kGroupSize + j; + float q = dl*qb[j]; + sumqx += wb[jj]*xb[jj]*q; + sumq2 += wb[jj]*q*q; + } + } + xb += Q::kBlockSize; + wb += Q::kBlockSize; + } + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; + } + } + + float id = d ? 1/d : 0.f; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto scales = all_scales + ibl*Q::kNblock; + for (int ib = 0; ib < Q::kNblock/2; ++ib) { + int ls1 = best_index_iq4nl(iq4k_values, id*scales[ib]); + int ls2 = best_index_iq4nl(iq4k_values, id*scales[ib + Q::kNblock/2]); + y[ibl].scales[ib] = ls1 | (ls2 << 4); + } + } + + *dptr = d; + if (!d) return; + + for (int iloop = 0; iloop < 1; ++iloop) { + + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + + auto qs = (uint16_t *)y[ibl].ql; + const float * xbl = x + ibl*Q::kSuperBlockSize; + + for (int ib = 0; ib < Q::kNblock; ++ib) { + const float * xb = xbl + Q::kBlockSize*ib; + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + int ls = iq4k_values[(y[ibl].scales[ib%(Q::kNblock/2)] >> 4*(ib/(Q::kNblock/2))) & 0xf]; + float dl = d*ls; + quantizer.find_best_match(dl, xb, weight, best_idx); + + auto prev_idx = all_idx + (ibl*Q::kSuperBlockSize + ib*Q::kBlockSize)/Q::kGroupSize; + + float mse1 = 0, mse2 = 0; + for (int ig = 0; ig < Q::kNg; ++ig) { + auto q1 = quantizer.values() + Q::kGroupSize*prev_idx[ig]; + auto q2 = quantizer.values() + Q::kGroupSize*best_idx[ig]; + for (int j = 0; j < Q::kGroupSize; ++j) { + int jj = ig*Q::kGroupSize + j; + float diff1 = xb[jj] - dl*q1[j]; + float diff2 = xb[jj] - dl*q2[j]; + mse1 += weight[jj]*diff1*diff1; + mse2 += weight[jj]*diff2*diff2; + } + } + if (mse1 < mse2) { + for (int ig = 0; ig < Q::kNg; ++ig) best_idx[ig] = prev_idx[ig]; + } else { + for (int ig = 0; ig < Q::kNg; ++ig) prev_idx[ig] = best_idx[ig]; + } + + for (int j = 0; j < Q::kNg; ++j) { + qs[j] = best_idx[j]; + auto xl = xb + Q::kGroupSize*j; + auto wl = weight + Q::kGroupSize*j; + auto ql = quantizer.values() + best_idx[j]*Q::kGroupSize; + for (int k = 0; k < Q::kGroupSize; ++k) { + float q = ql[k]*ls; + sumqx += wl[k]*xl[k]*q; + sumq2 += wl[k]*q*q; + } + } + qs += Q::kNg; + } + } + if (sumq2 > 0) { + d = sumqx/sumq2; + *dptr = d; + if (!d) return; + } else { + break; + } + + if (false) { + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * xbl = x + ibl*Q::kSuperBlockSize; + auto scales = all_scales + ibl*Q::kNblock; + auto qs = (uint16_t *)y[ibl].ql; + for (int ib = 0; ib < Q::kNblock; ++ib) { + const float * xb = xbl + Q::kBlockSize*ib; + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + for (int j = 0; j < Q::kNg; ++j) best_idx[j] = qs[ib*Q::kNg+j]; + auto pair = quantizer.find_best_scale(xb, weight, best_idx); + scales[ib] = pair.first; + } + } + float id = d ? 1/d : 0.f; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto scales = all_scales + ibl*Q::kNblock; + for (int ib = 0; ib < Q::kNblock/2; ++ib) { + int ls1 = best_index_iq4nl(iq4k_values, id*scales[ib]); + int ls2 = best_index_iq4nl(iq4k_values, id*scales[ib + Q::kNblock/2]); + y[ibl].scales[ib] = ls1 | (ls2 << 4); + } + } + } + + } + +} +} + +void quantize_row_iq2_kt_ref(const float * GGML_RESTRICT x, block_iq2_kt * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq2_kt(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq2_kt(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_iq2_kt * y = (block_iq2_kt *)vy; + quantize_row_iq2_kt_ref(x, y, k); +} + +size_t quantize_iq2_kt(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ2_KT, n_per_row); + std::vector scales(n_per_row/QuantizerIQ2KT::kBlockSize); + std::vector weights(n_per_row); + std::vector idx(n_per_row/QuantizerIQ2KT::kGroupSize); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrows; ++row) { + quantize_row_iq2_kt_impl(src, (void *)qrow, n_per_row, imatrix, scales.data(), weights.data(), idx.data()); + src += n_per_row; + qrow += row_size; + } + return nrows * row_size; +} + +void dequantize_row_iq2_kt(const block_iq2_kt * x, float * y, int64_t k) { + assert(k % QuantizerIQ2KT::kSuperBlockSize == 0); +#ifdef __AVX2__ + //if (iqk_dequantize_ktquants(GGML_TYPE_IQ2_KT, k, x, 0, y, 0, 1)) return; +#endif + const int nb = k / QuantizerIQ2KT::kSuperBlockSize; + const float * dptr = (const float *)x; + const float d = *dptr * QuantizerIQ2KT::kScale; + x = (const block_iq2_kt *)(dptr + 1); + auto& deq = iq2kt_quantizer(); + for (int ibl = 0; ibl < nb; ++ibl) { + auto yl = y + ibl*QuantizerIQ2KT::kSuperBlockSize; + auto yh = yl + QuantizerIQ2KT::kSuperBlockSize/2; + const uint16_t * ql = (const uint16_t *)x[ibl].ql; + const uint16_t * qh = ql + QuantizerIQ2KT::kNg*QuantizerIQ2KT::kNblock/2; + for (int ib = 0; ib < QuantizerIQ2KT::kNblock/2; ++ib) { + float sl = d * iq4k_values[x[ibl].scales[ib] & 0xf]; + float sh = d * iq4k_values[x[ibl].scales[ib] >> 4]; + for (int ig = 0; ig < QuantizerIQ2KT::kNg; ++ig) { + deq.set_values(ql[ig], yl, sl); + deq.set_values(qh[ig], yh, sh); + yl += QuantizerIQ2KT::kGroupSize; + yh += QuantizerIQ2KT::kGroupSize; + } + ql += QuantizerIQ2KT::kNg; + qh += QuantizerIQ2KT::kNg; + } + } +} + +void vec_dot_iq2_kt_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ2_KT, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + +} + +namespace { + +using QuantizerIQ3KT = QuantizerIQKT<32, 8, 16, true, true>; +const QuantizerIQ3KT& iq3kt_quantizer() { + static std::mutex mutex; + std::lock_guard lock(mutex); + static std::unique_ptr quantizer; + if (!quantizer) quantizer = std::make_unique(256, 8); + return *quantizer; +} + +void quantize_row_iq3_kt_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, float * all_scales, + float * all_weights, float * qtmp) { + + constexpr float kSigmaScale = 2.0f; + constexpr float kStep = 8.0f; + + using Q = QuantizerIQ3KT; + + static_assert(Q::kNumVal%8 == 0); + + constexpr int kNumGroups = Q::kSuperBlockSize/Q::kGroupSize; + + float * dptr = (float *)vy; + + block_iq3_kt * y = (block_iq3_kt *)(dptr + 1); + + int best_idx[2*Q::kNg]; + + auto& quantizer = iq3kt_quantizer(); + + int nblock = n_per_row / Q::kSuperBlockSize; + + float amax_row = 0; + for (int j = 0; j < n_per_row; ++j) amax_row = std::max(amax_row, std::abs(x[j])); + if (!amax_row) { + *dptr = 0.f; + std::memset(y, 0, nblock*sizeof(block_iq3_kt)); + return; + } + + Q::set_weights(kSigmaScale, nblock, x, quant_weights, all_weights); + + float amax_scale = 0, max_scale = 0; + + float xaux[Q::kBlockSize]; + + for (int ibl = 0; ibl < nblock; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq3_kt)); + + auto scales = all_scales + ibl*Q::kNblock; + auto xbl = x + ibl*Q::kSuperBlockSize; + + for (int ib = 0; ib < Q::kNblock; ++ib) { + const float * xb = xbl + Q::kBlockSize*ib; + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + float amax = 0; + for (int j = 0; j < Q::kBlockSize; ++j) { + float ax = std::abs(xb[j]); + xaux[j] = ax; + amax = std::max(amax, ax); + } + if (amax < 1e-16f) { + scales[ib] = 0.0f; + continue; + } + + //quantizer.find_best_match(amax/96.f, xaux, weight, best_idx+Q::kNg); + //scales[ib] = quantizer.find_best_scale(xaux, weight, best_idx+Q::kNg).first; + + float scale_0 = std::max(84.f, 123.f*amax/amax_row); + //float scale_0 = std::max(64.f, 123.f*amax/amax_row); + float best = 0; + bool found_solution = false; + for (int itry = -3; itry <= 3; ++itry) { + quantizer.find_best_match(amax/(scale_0 + kStep*itry), xaux, weight, best_idx); + auto [d, score] = quantizer.find_best_scale(xaux, weight, best_idx); + if (score > best) { + best = score; + found_solution = true; + scales[ib] = d; + std::memcpy(best_idx+Q::kNg, best_idx, Q::kNg*sizeof(int)); + } + } + if (!found_solution) { + fprintf(stderr, "======================= %s: failed to find solution for a block\n", __func__); + fprintf(stderr, "Model weights and importances:\n"); + for (int j = 0; j < Q::kBlockSize; ++j) { + fprintf(stderr, "%2d %g %g\n", j, xaux[j], weight[j]); + } + GGML_ASSERT(false); + } + + auto xt = qtmp + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + for (int ig = 0; ig < Q::kNg; ++ig) { + auto q = quantizer.values() + Q::kGroupSize*best_idx[Q::kNg+ig]; + for (int j = 0; j < Q::kGroupSize; ++j) *xt++ = q[j]; + } + + float abs_scale = std::abs(scales[ib]); + if (abs_scale > amax_scale) { + amax_scale = abs_scale; + max_scale = scales[ib]; + } + } + + } + + GGML_ASSERT(max_scale >= 0); + float d = max_scale/15; + float best = 0; + for (int itry = -9; itry <= 9; ++itry) { + float id = (itry*0.2f + 15)/max_scale; + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + const float * xb = x + ibl*Q::kSuperBlockSize; + const float * qb = qtmp + ibl*Q::kSuperBlockSize; + const float * wb = all_weights + ibl*Q::kSuperBlockSize; + auto scales = all_scales + ibl*Q::kNblock; + for (int ib = 0; ib < Q::kNblock; ++ib) { + int ls = nearest_int(id*scales[ib]); + ls = std::max(0, std::min(15, ls)); + float dl = ls; + for (int j = 0; j < Q::kBlockSize; ++j) { + float q = dl*qb[j]; + sumqx += wb[j]*std::abs(xb[j])*q; + sumq2 += wb[j]*q*q; + } + xb += Q::kBlockSize; + wb += Q::kBlockSize; + qb += Q::kBlockSize; + } + } + if (sumq2 > 0 && sumqx*sumqx > best*sumq2) { + d = sumqx/sumq2; best = d*sumqx; + } + } + + float id = d ? 1/d : 0.f; + for (int ibl = 0; ibl < nblock; ++ibl) { + auto scales = all_scales + ibl*Q::kNblock; + for (int ib = 0; ib < Q::kNblock/2; ++ib) { + int ls1 = nearest_int(id*scales[ib]); + int ls2 = nearest_int(id*scales[ib + Q::kNblock/2]); + ls1 = std::max(0, std::min(15, ls1)); + ls2 = std::max(0, std::min(15, ls2)); + y[ibl].scales[ib] = ls1 | (ls2 << 4); + } + } + + *dptr = d; + + for (int iloop = 0; iloop < 1; ++iloop) { + + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + + uint16_t * ql = (uint16_t *)y[ibl].ql; + + std::memset(y[ibl].qh, 0, kNumGroups/2); + const float * xbl = x + ibl*Q::kSuperBlockSize; + + for (int ib = 0; ib < Q::kNblock; ++ib) { + const float * xb = xbl + Q::kBlockSize*ib; + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + for (int j = 0; j < Q::kBlockSize; ++j) { + xaux[j] = std::abs(xb[j]); + if (xb[j] < 0) y[ibl].qh[j] |= (1 << ib); + } + int ls = (y[ibl].scales[ib%(Q::kNblock/2)] >> 4*(ib/(Q::kNblock/2))) & 0xf; + float dl = d*ls; + quantizer.find_best_match(dl, xaux, weight, best_idx); + + for (int j = 0; j < Q::kNg; ++j) { + ql[ib*Q::kNg+j] = best_idx[j]; + auto xl = xaux + Q::kGroupSize*j; + auto wl = weight + Q::kGroupSize*j; + auto ql = quantizer.values() + best_idx[j]*Q::kGroupSize; + for (int k = 0; k < Q::kGroupSize; ++k) { + float q = ql[k]*ls; + sumqx += wl[k]*xl[k]*q; + sumq2 += wl[k]*q*q; + } + } + } + } + if (sumq2 > 0) { + d = sumqx/sumq2; + *dptr = d; + if (!d) break; + } else { + break; + } + } +} +} + +void quantize_row_iq3_kt_ref(const float * x, block_iq3_kt * y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq3_kt(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq3_kt(const float * x, void * vy, int64_t k) { + assert(k % QK_K == 0); + block_iq3_kt * y = (block_iq3_kt *)vy; + quantize_row_iq3_kt_ref(x, y, k); +} + +size_t quantize_iq3_kt(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ3_KT, n_per_row); + std::vector scales(n_per_row/QuantizerIQ3KT::kBlockSize); + std::vector weights(n_per_row), xtmp(n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrows; ++row) { + quantize_row_iq3_kt_impl(src, (void *)qrow, n_per_row, imatrix, scales.data(), weights.data(), xtmp.data()); + src += n_per_row; + qrow += row_size; + } + return nrows * row_size; +} + +void dequantize_row_iq3_kt(const block_iq3_kt * x, float * y, int64_t k) { +#ifdef __AVX2__ + //if (iqk_dequantize_ktquants(GGML_TYPE_IQ3_KT, k, x, 0, y, 0, 1)) return; +#endif + using Q = QuantizerIQ3KT; + constexpr int kNumGroups = Q::kSuperBlockSize/Q::kGroupSize; + assert(k % Q::kSuperBlockSize == 0); + const int nb = k / Q::kSuperBlockSize; + const float * dptr = (const float *)x; + const float d = *dptr * Q::kScale; + x = (const block_iq3_kt *)(dptr + 1); + auto& deq = iq3kt_quantizer(); + for (int ibl = 0; ibl < nb; ++ibl) { + auto yl = y + ibl*Q::kSuperBlockSize; + auto yh = yl + Q::kSuperBlockSize/2; + auto qll = (const uint16_t *)x[ibl].ql; + auto qlh = qll + kNumGroups/2; + int jj = 0; + for (int ib = 0; ib < Q::kNblock/2; ++ib) { + float sl = d * (x[ibl].scales[ib] & 0xf); + float sh = d * (x[ibl].scales[ib] >> 4); + uint8_t l_mask = 1 << ib; + uint8_t h_mask = l_mask << (Q::kNblock/2); + for (int ig = 0; ig < Q::kNg; ++ig) { + deq.set_values(qll[jj], yl, sl); + deq.set_values(qlh[jj], yh, sh); + for (int j = 0; j < Q::kGroupSize; ++j) { + if (x[ibl].qh[ig*Q::kGroupSize+j] & l_mask) yl[j] = -yl[j]; + if (x[ibl].qh[ig*Q::kGroupSize+j] & h_mask) yh[j] = -yh[j]; + } + yl += Q::kGroupSize; + yh += Q::kGroupSize; + ++jj; + } + } + } +} + +void vec_dot_iq3_kt_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ3_KT, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + +} + +// ======================================== iq4_kt + +namespace{ + +using QuantizerIQ4KT = QuantizerIQKT<32, 4, 15, false, true>; + +const QuantizerIQ4KT& iq4kt_quantizer(bool with_offset = false) { + static std::mutex mutex; + std::lock_guard lock(mutex); + static std::unique_ptr quantizer1; + static std::unique_ptr quantizer2; + if (with_offset) { + if (!quantizer2) quantizer2 = std::make_unique(625, 6, 4096+32768); + return *quantizer2; + } + if (!quantizer1) quantizer1 = std::make_unique(625, 6, 4096); + return *quantizer1; +} + +const QuantizerIQ4KT& iq4kt_dequantizer() { + static std::mutex mutex; + std::lock_guard lock(mutex); + static std::unique_ptr dequantizer; + if (!dequantizer) dequantizer = std::make_unique(0, 0, 4096); + return *dequantizer; +} + +void quantize_row_iq4_kt_impl(const float * x, void * vy, int n_per_row, const float * quant_weights, float * all_scales, float * all_weights) { + + constexpr float kSigmaScale = 2.0f; + constexpr int kNtry = 2; + using Q = QuantizerIQ4KT; + + static_assert(Q::kNumVal%8 == 0); + + float * dptr = (float *)vy; + + block_iq4_kt * y = (block_iq4_kt *)(dptr + 1); + + auto& quantizer1 = iq4kt_quantizer(); + auto& quantizer2 = iq4kt_quantizer(true); + + int nblock = n_per_row / Q::kSuperBlockSize; + + Q::set_weights(kSigmaScale, nblock, x, quant_weights, all_weights); + + float amax_row = 0; + for (int j = 0; j < n_per_row; ++j) { + amax_row = std::max(amax_row, std::abs(x[j])); + } + if (!amax_row) { + dptr[0] = 0.f; + std::memset(y, 0, nblock*sizeof(block_iq4_kt)); + return; + } + + int best_idx[2*Q::kNg]; + float xaux[Q::kBlockSize]; + + float amax_scale = 0, max_scale = 0; + + for (int ibl = 0; ibl < nblock; ++ibl) { + + memset(&y[ibl], 0, sizeof(block_iq4_kt)); + + const float * xbl = x + ibl*Q::kSuperBlockSize; + auto scales = all_scales + ibl*Q::kNblock; + + for (int ib = 0; ib < Q::kNblock; ++ib) { + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + float amax = 0; + for (int j = 0; j < Q::kBlockSize; ++j) { + xaux[j] = xbl[ib*Q::kBlockSize+j]; + float ax = std::abs(xaux[j]); + amax = std::max(amax, ax); + } + if (amax < 1e-16f) { + scales[ib] = 0; + continue; + } + float best = 0; + float scale_0 = std::max(90.f, 124.f*amax/amax_row); + for (int itry = -kNtry; itry <= kNtry; ++itry) { + quantizer1.find_best_match( amax/(8.f*itry + scale_0), xaux, weight, best_idx); + auto [dp, score_p] = quantizer1.find_best_scale(xaux, weight, best_idx); + if (score_p > best) { + best = score_p; scales[ib] = dp; + } + quantizer1.find_best_match(-amax/(8.f*itry + scale_0), xaux, weight, best_idx); + auto [dm, score_m] = quantizer1.find_best_scale(xaux, weight, best_idx); + if (score_m > best) { + best = score_m; scales[ib] = dm; + } + } + + quantizer2.find_best_match(scales[ib], xaux, weight, best_idx); + auto [d, score] = quantizer2.find_best_scale(xaux, weight, best_idx); + if (score > best) { + scales[ib] = d; + y[ibl].qs[ib] = 1; + } + bool with_offset = false; + for (int itry = -kNtry; itry <= kNtry; ++itry) { + quantizer2.find_best_match( amax/(8.f*itry + scale_0), xaux, weight, best_idx); + auto [dp, score_p] = quantizer2.find_best_scale(xaux, weight, best_idx); + if (score_p > best) { + best = score_p; scales[ib] = dp; with_offset = true; + } + quantizer2.find_best_match(-amax/(8.f*itry + scale_0), xaux, weight, best_idx); + auto [dm, score_m] = quantizer2.find_best_scale(xaux, weight, best_idx); + if (score_m > best) { + best = score_m; scales[ib] = dm; with_offset = true; + } + } + if (with_offset) y[ibl].qs[ib] = 1; + + float abs_scale = std::abs(scales[ib]); + if (abs_scale > amax_scale) { + amax_scale = abs_scale; + max_scale = scales[ib]; + } + } + + } + + float d = -max_scale/64; + + dptr[0] = d; + if (!d) return; + + constexpr int kNumGroups = Q::kSuperBlockSize/Q::kGroupSize; + + for (int iloop = 0; iloop < 1; ++iloop) { + + const float id = 1/d; + + float sumqx = 0, sumq2 = 0; + for (int ibl = 0; ibl < nblock; ++ibl) { + + // high 3 bits + scales + // each block of 32 needs 8 x 3 (high bits) + 1 x 8 (scale) = 32 bits = 1 x uint32_t + // we have 8 blocks + auto shb = y[ibl].qs; // high 3 bits + scales + auto ql = (uint8_t *)(shb + Q::kNblock); + auto qh = ql + kNumGroups; + std::memset(qh, 0, kNumGroups/2); + const float * xbl = x + ibl*Q::kSuperBlockSize; + auto scales = all_scales + ibl*Q::kNblock; + + for (int ib = 0; ib < Q::kNblock; ++ib) { + auto& quantizer = y[ibl].qs[ib] & 1 ? quantizer2 : quantizer1; + const float * weight = all_weights + ibl*Q::kSuperBlockSize + ib*Q::kBlockSize; + for (int j = 0; j < Q::kBlockSize; ++j) xaux[j] = xbl[ib*Q::kBlockSize+j]; + int ls = nearest_int(id*scales[ib]); + ls = std::min(ls, 63); + *(uint8_t *)(shb + ib) = ((ls + 64) << 1) | (shb[ib] & 1); + float dl = d*ls; + quantizer.find_best_match(dl, xaux, weight, best_idx); + + for (int j = 0; j < Q::kNg; ++j) { + shb[ib] |= ((best_idx[j] >> 12) << (8 + 3*j)); + ql[Q::kNg*ib + j] = best_idx[j] & 255; + qh[(Q::kNg*ib + j)%(kNumGroups/2)] |= ((best_idx[j] >> 8) & 0xf) << 4*((Q::kNg*ib + j)/(kNumGroups/2)); + auto xl = xaux + Q::kGroupSize*j; + auto wl = weight + Q::kGroupSize*j; + auto ql = quantizer.values() + Q::kGroupSize*best_idx[j]; + for (int k = 0; k < Q::kGroupSize; ++k) { + float q = ql[k]*ls; + sumqx += wl[k]*xl[k]*q; + sumq2 += wl[k]*q*q; + } + } + } + } + if (sumq2 > 0) { + d = sumqx/sumq2; + dptr[0] = d; + if (!d) break; + } else { + break; + } + } +} +} + +void quantize_row_iq4_kt_ref(const float * GGML_RESTRICT x, block_iq4_kt * GGML_RESTRICT y, int64_t k) { + assert(k % QK_K == 0); + quantize_iq4_kt(x, (void *)y, 1, k, nullptr, nullptr); +} + +void quantize_row_iq4_kt(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) { + assert(k % QK_K == 0); + block_iq4_kt * y = (block_iq4_kt *)vy; + quantize_row_iq4_kt_ref(x, y, k); +} + +size_t quantize_iq4_kt(const float * src, void * dst, int64_t nrows, int64_t n_per_row, const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row%QK_K == 0); + auto row_size = ggml_row_size(GGML_TYPE_IQ4_KT, n_per_row); + std::vector scales(n_per_row/QuantizerIQ4KT::kBlockSize); + std::vector weights(n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrows; ++row) { + quantize_row_iq4_kt_impl(src, (void *)qrow, n_per_row, imatrix, scales.data(), weights.data()); + src += n_per_row; + qrow += row_size; + } + return nrows * row_size; +} + +void dequantize_row_iq4_kt(const block_iq4_kt * x, float * y, int64_t k) { +#ifdef __AVX2__ + //if (iqk_dequantize_ktquants(GGML_TYPE_IQ4_KT, k, x, 0, y, 0, 1)) return; +#endif + using Q = QuantizerIQ4KT; + assert(k % Q::kSuperBlockSize == 0); + constexpr int kNumGroups = Q::kSuperBlockSize/Q::kGroupSize; + const int nb = k / Q::kSuperBlockSize; + const float * dptr = (const float *)x; + const float d = dptr[0] * Q::kScale; + x = (const block_iq4_kt *)(dptr + 1); + auto& deq = iq4kt_dequantizer(); + for (int ibl = 0; ibl < nb; ++ibl) { + auto shb = x[ibl].qs; + auto ql = (const uint8_t *)(shb + Q::kNblock); + auto qh = ql + kNumGroups; + for (int ib = 0; ib < Q::kNblock; ++ib) { + int offset = shb[ib] & 1 ? 32768 + 4096 : 4096; + int ls = int((shb[ib] & 0xff) >> 1) - 64; + float sl = d * ls; + for (int ig = 0; ig < Q::kNg; ++ig) { + int jj = ib*Q::kNg+ig; + uint16_t idx = ql[jj] | ((qh[jj%(kNumGroups/2)] << (8 - 4*(jj/(kNumGroups/2)))) & 0xf00) | (((shb[ib] >> (8 + 3*ig)) & 7) << 12); + deq.set_values(idx, y, sl, offset); + y += Q::kGroupSize; + } + } + } +} + +void vec_dot_iq4_kt_q8_k(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK_K == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + +#if GGML_USE_IQK_MULMAT + if (iqk_mul_mat(1, 1, n, GGML_TYPE_IQ4_KT, vx, 0, GGML_TYPE_Q8_K, vy, 0, s, 0, 0, 1)) { + return; + } +#endif + +} + +void quantize_row_q1_0_g128_ref(const float * x, block_q1_0_g128 * y, int64_t k) { + quantize_row_q1_0_g128(x, y, k); +} + +void quantize_row_q1_0_g128(const float * x, void * vy, int64_t k) { + assert(k % QK1_0_G128 == 0); + int nb = k / QK1_0_G128; + auto y = (block_q1_0_g128 *)vy; + for (int ib = 0; ib < nb; ++ib) { + float sum = 0; + for (int j = 0; j < QK1_0_G128; ++j) sum += std::abs(x[j]); + float d = sum / QK1_0_G128; + y[ib].d = GGML_FP32_TO_FP16(d); + std::memset(y[ib].qs, 0, QK1_0_G128/8); + for (int j = 0; j < QK1_0_G128; ++j) { + if (x[j] >= 0.0f) { + y[ib].qs[j / 8] |= (1 << (j % 8)); + } + } + x += QK1_0_G128; + } +} + +size_t quantize_q1_0_g128(const float * src, void * dst, int64_t nrows, int64_t n_per_row, [[maybe_unused]] const float * imatrix, + [[maybe_unused]] const quantize_user_data * user_data) { + GGML_ASSERT(n_per_row % QK1_0_G128 == 0); + int64_t ntot = nrows * n_per_row; + quantize_row_q1_0_g128(src, dst, ntot); + int64_t nblock = ntot / QK1_0_G128; + return nblock * sizeof(block_q1_0_g128); +} + +void dequantize_row_q1_0_g128(const block_q1_0_g128 * x, float * y, int64_t k) { + assert(k % QK1_0_G128 == 0); + constexpr uint8_t k_mask[8] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80}; + int nb = k / QK1_0_G128; + for (int ib = 0; ib < nb; ++ib) { + float d = GGML_FP16_TO_FP32(x[ib].d); + for (int i = 0; i < QK1_0_G128/8; ++i) { + for (int j = 0; j < 8; ++j) { + *y++ = x[ib].qs[i] & k_mask[j] ? d : -d; + } + } + } +} + +void vec_dot_q1_0_g128_q8_0(int n, float * s, size_t bs, const void * vx, size_t bx, const void * vy, size_t by, int nrc) { + assert(n % QK1_0_G128 == 0); + assert(nrc == 1); + GGML_UNUSED(nrc); + GGML_UNUSED(bx); + GGML_UNUSED(by); + GGML_UNUSED(bs); + int nb = n / QK1_0_G128; + + constexpr uint8_t k_mask[8] = {0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80}; + + constexpr int n4 = QK1_0_G128 / QK8_0; + + auto x = (const block_q1_0_g128 *)vx; + auto y = (const block_q8_0_x4 *)vy; + int16_t sumi[QK1_0_G128/8]; + float sumf = 0; + for (int ib = 0; ib < nb; ++ib) { + auto dx = GGML_FP16_TO_FP32(x[ib].d); + auto qx = x[ib].qs; + auto qy = y[ib].qs; + for (int k = 0; k < QK1_0_G128/8; ++k) { + uint8_t bits = qx[k]; + int16_t s = 0; + for (int j = 0; j < 8; ++j) { + s += (bits & k_mask[j] ? qy[j] : -qy[j]); + } + qy += 8; + sumi[k] = s; + } + auto s = sumi; + for (int k = 0; k < n4; ++k) { + float dy = GGML_FP16_TO_FP32(y[ib].d[k]); + sumf += dx*dy*(s[0] + s[1] + s[2] + s[3]); + s += 4; + } + } + *s = sumf; +} + +namespace { +template +inline int check_row_for_blocks_256_fp16(int nblock, const Block * x) { + int nbad = 0; + for (int ib = 0; ib < nblock; ++ib) { + float d = GGML_FP16_TO_FP32(x[ib].d); + if (isnan(d)) ++nbad; + } + return nbad; +} +template +bool check_tensor_for_blocks_256_fp16(const ggml_tensor * tensor) { + int nblock = tensor->ne[0]/QK_K; + int nbad = 0; + for (int row = 0; row < ggml_nrows(tensor); ++row) { + auto x = (const Block *)((const char *)tensor->data + tensor->nb[1]*row); + nbad += check_row_for_blocks_256_fp16(nblock, x); + } + if (nbad > 0) { + fprintf(stderr, "%s: found %d NaN block scales out of %g blocks in tensor %s\n", __func__, + nbad, 1.*ggml_nrows(tensor)*nblock, tensor->name); + if (tensor->ne[2] > 1) { + int nb = tensor->ne[0]/QK_K; + for (int64_t i02 = 0; i02 < tensor->ne[2]; ++i02) { + int nbad_expert = 0; + auto xex = (const char *)((const char *)tensor->data + i02*tensor->nb[2]); + for (int64_t i01 = 0; i01 < tensor->ne[1]; ++i01) { + auto xr = (const Block *)(xex + i01*tensor->nb[1]); + nbad_expert += check_row_for_blocks_256_fp16(nb, xr); + } + if (nbad_expert > 0) fprintf(stderr," there are %d NaN block scales for expert %g\n", nbad_expert, 1.*i02); + } + } + return false; + } + return true; +} +template +inline int check_row_for_blocks_256_fp16(int nblock, const Block * x, int nr) { + int nbad = 0; + for (int ib = 0; ib < nblock; ++ib) { + for (int j = 0; j < nr; ++j) { + if (!isfinite(GGML_FP16_TO_FP32(x[ib].d[j]))) ++nbad; + } + } + return nbad; +} +template +bool check_tensor_for_blocks_256_fp16_repacked(const ggml_tensor * tensor) { + int nblock = tensor->ne[0]/QK_K; + int nbad = 0; + for (int row = 0; row < ggml_nrows(tensor); row += nr) { + auto x = (const Block *)((const char *)tensor->data + tensor->nb[1]*row); + nbad += check_row_for_blocks_256_fp16(nblock, x, nr); + } + if (nbad > 0) { + fprintf(stderr, "%s: found %d NaN block scales out of %g blocks in tensor %s\n", __func__, + nbad, 1.*ggml_nrows(tensor)*nblock, tensor->name); + if (tensor->ne[2] > 1) { + int nb = tensor->ne[0]/QK_K; + for (int64_t i02 = 0; i02 < tensor->ne[2]; ++i02) { + int nbad_expert = 0; + auto xex = (const char *)((const char *)tensor->data + i02*tensor->nb[2]); + for (int64_t i01 = 0; i01 < tensor->ne[1]; i01 += nr) { + auto xr = (const Block *)(xex + i01*tensor->nb[1]); + nbad_expert += check_row_for_blocks_256_fp16(nb, xr, nr); + } + if (nbad_expert > 0) fprintf(stderr," there are %d NaN block scales for expert %g\n", nbad_expert, 1.*i02); + } + } + return false; + } + return true; +} +struct F32Scale { + static inline int check_row(const char * data) { + float d = *(const float *)data; + return isfinite(d) ? 0 : 1; + } +}; +struct F16Scale { + static inline int check_row(const char * data) { + float d = GGML_FP16_TO_FP32(*(const ggml_half *)data); + return isfinite(d) ? 0 : 1; + } +}; +template +struct F32ScaleRX { + static inline int check_row(const char * data) { + auto d = (const float *)data; + int nbad = 0; + for (int i = 0; i < nr; ++i) { + if (!isfinite(d[i])) ++nbad; + } + return nbad; + } +}; +template +struct F16ScaleRX { + static inline int check_row(const char * data) { + auto d = (const ggml_half *)data; + int nbad = 0; + for (int i = 0; i < nr; ++i) { + if (!isfinite(GGML_FP16_TO_FP32(d[i]))) ++nbad; + } + return nbad; + } +}; +template +bool check_tensor_row_scales(const ggml_tensor * tensor) { + auto row_size = ggml_row_size(tensor->type, tensor->ne[0]); + int num_rows = ggml_nrows(tensor); + auto data = (const char *)tensor->data; + int nbad = 0; + for (int row = 0; row < num_rows; ++row) { + nbad += RS::check_row(data); + data += row_size; + } + if (nbad > 0) { + fprintf(stderr, "%s: found %d NaN row scales out of %d rows in tensor %s\n", __func__, + nbad, num_rows, tensor->name); + return false; + } + return true; +} +} + +bool iqk_validate_tensor(const ggml_tensor * tensor) { + if (!tensor) return true; + if (!ggml_is_contiguous(tensor)) return true; + + switch (tensor->type) { + case GGML_TYPE_IQ2_K: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ3_K: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ4_K: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ5_K: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ6_K: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ2_XXS: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ2_XS: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ2_S: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ3_XXS: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ3_S: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ4_XS: return check_tensor_for_blocks_256_fp16(tensor); + case GGML_TYPE_IQ2_K_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ3_K_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ4_K_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ5_K_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ2_XXS_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ2_XS_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ2_S_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ3_XXS_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ3_S_R4: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ4_XS_R8: return check_tensor_for_blocks_256_fp16_repacked(tensor); + case GGML_TYPE_IQ2_BN: + case GGML_TYPE_IQ4_KSS: + case GGML_TYPE_IQ4_KS: + case GGML_TYPE_IQ5_KS: return check_tensor_row_scales(tensor); + case GGML_TYPE_IQ2_BN_R4: + case GGML_TYPE_IQ4_KS_R4: + case GGML_TYPE_IQ5_KS_R4: return check_tensor_row_scales>(tensor); + case GGML_TYPE_IQ1_BN: + case GGML_TYPE_IQ2_KS: + case GGML_TYPE_IQ2_KL: + case GGML_TYPE_IQ3_KS: return check_tensor_row_scales(tensor); + case GGML_TYPE_IQ1_S_R4: + case GGML_TYPE_IQ1_M_R4: return check_tensor_row_scales>(tensor); + + default: break; + } + return true; +} diff --git a/llama.cpp/ggml/src/iqk/iqk_quantize.h b/llama.cpp/ggml/src/iqk/iqk_quantize.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_quantize.h @@ -0,0 +1,385 @@ +// opencoti F5-opt W2 (#290): the mxfp4-family quant entry points (quantize_mxfp4 +// et al.) were renamed iqk_*_mxfp4 — ik_llama gives them a 6-arg (quantize_user_data) +// signature that collides with llamafile's own 5-arg quantize_mxfp4 in ggml-quants.{h,c}. +// We never call ik_llama's mxfp4 path; the rename keeps both vintages linkable. Re-sync: reapply. +// +// Copyright (C) 2024-2025 Iwan Kawrakow +// MIT license +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include +#include + +#define GGML_COMMON_DECL_C +#include "ggml-common.h" + +#ifdef __cplusplus +#define GGML_RESTRICT +extern "C" { +#else +#define GGML_RESTRICT restrict +#endif + +// opencoti F5-opt W2 (#290): ik_llama defines this in its own ggml/include/ggml.h +// (a superset of llamafile's ggml.h, which we must not edit). iqk_quantize.cpp +// dereferences user_data->slow_iq2_ks, so the struct must be COMPLETE here, not +// merely forward-declared. Body copied verbatim from ik_llama ggml.h. +#ifdef __cplusplus +struct quantize_user_data { + bool symmetric_q4_0; + bool slow_iq2_ks; +}; +#else +struct quantize_user_data; +#endif + +void quantize_row_iq2_k_ref(const float * GGML_RESTRICT x, block_iq2_k * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_k(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_k(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq2_k(const block_iq2_k * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq2_k_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq3_k_ref(const float * GGML_RESTRICT x, block_iq3_k * GGML_RESTRICT y, int64_t k); +void quantize_row_iq3_k(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq3_k(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq3_k(const block_iq3_k * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq3_k_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq3_ks_ref(const float * GGML_RESTRICT x, block_iq3_ks * GGML_RESTRICT y, int64_t k); +void quantize_row_iq3_ks(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq3_ks(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq3_ks(const block_iq3_ks * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq3_ks_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq4_k_ref(const float * GGML_RESTRICT x, block_iq4_k * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_k(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq4_k(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq4_k(const block_iq4_k * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq4_k_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq5_k_ref(const float * GGML_RESTRICT x, block_iq5_k * GGML_RESTRICT y, int64_t k); +void quantize_row_iq5_k(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq5_k(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq5_k(const block_iq5_k * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq5_k_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq6_k_ref(const float * GGML_RESTRICT x, block_iq6_k * GGML_RESTRICT y, int64_t k); +void quantize_row_iq6_k(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq6_k(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq6_k(const block_iq6_k * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq6_k_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq4_ks_ref(const float * GGML_RESTRICT x, block_iq4_ks * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_ks(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq4_ks(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq4_ks(const block_iq4_ks * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq4_ks_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq4_kss_ref(const float * GGML_RESTRICT x, block_iq4_kss * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_kss(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq4_kss(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq4_kss(const block_iq4_kss * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq4_kss_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void iqk_quantize_row_mxfp4_ref(const float * GGML_RESTRICT x, block_mxfp4 * GGML_RESTRICT y, int64_t k); +void iqk_quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t iqk_quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void iqk_dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_mxfp4_q8_0_x4(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_ks_ref(const float * GGML_RESTRICT x, block_iq2_ks * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_ks(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_ks(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq2_ks(const block_iq2_ks * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq2_ks_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_kl_ref(const float * GGML_RESTRICT x, block_iq2_kl * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_kl(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_kl(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq2_kl(const block_iq2_kl * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq2_kl_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq1_kt_ref(const float * GGML_RESTRICT x, block_iq1_kt * GGML_RESTRICT y, int64_t k); +void quantize_row_iq1_kt(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq1_kt(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq1_kt(const block_iq1_kt * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq1_kt_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_kt_ref(const float * GGML_RESTRICT x, block_iq2_kt * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_kt(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_kt(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq2_kt(const block_iq2_kt * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq2_kt_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq3_kt_ref(const float * GGML_RESTRICT x, block_iq3_kt * GGML_RESTRICT y, int64_t k); +void quantize_row_iq3_kt(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq3_kt(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq3_kt(const block_iq3_kt * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq3_kt_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq4_kt_ref(const float * GGML_RESTRICT x, block_iq4_kt * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_kt(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq4_kt(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq4_kt(const block_iq4_kt * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq4_kt_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq5_ks_ref(const float * GGML_RESTRICT x, block_iq5_ks * GGML_RESTRICT y, int64_t k); +void quantize_row_iq5_ks(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq5_ks(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq5_ks(const block_iq5_ks * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq5_ks_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq4_nl_r4_ref(const float * GGML_RESTRICT x, block_iq4_nl_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_nl_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq4_nl_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq4_nl_r4(const block_iq4_nl_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq4_nl_r4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q4_0_r8_ref(const float * GGML_RESTRICT x, block_iq4_nl_r8 * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_0_r8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q4_0_r8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q4_0_r8(const block_iq4_nl_r8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q4_0_r8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q8_0_r8_ref(const float * GGML_RESTRICT x, block_q8_0_r8 * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_0_r8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q8_0_r8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q8_0_r8(const block_q8_0_r8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q8_0_r8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q5_0_r4_ref(const float * GGML_RESTRICT x, block_q5_0_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_0_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q5_0_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q5_0_r4(const block_q5_0_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q5_0_r4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q6_0_r4_ref(const float * GGML_RESTRICT x, block_q6_0_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_q6_0_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q6_0_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q6_0_r4(const block_q6_0_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q6_0_r4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq4_xs_r8_ref(const float * GGML_RESTRICT x, block_iq4_xs_r8 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_xs_r8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq4_xs_r8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq4_xs_r8(const block_iq4_xs_r8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq4_xs_r8_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_bn_ref (const float * GGML_RESTRICT x, block_iq2_bn * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_bn (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void dequantize_row_iq2_bn (const block_iq2_bn * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_bn (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void vec_dot_iq2_bn_q8_K64(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_bn_r4_ref (const float * GGML_RESTRICT x, block_iq2_bn * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_bn_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void dequantize_row_iq2_bn_r4(const block_iq2_bn * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_bn_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void vec_dot_iq2_bn_r4_q8_K64(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q3_k_r4_ref(const float * GGML_RESTRICT x, block_q3_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_q3_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q3_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q3_k_r4(const block_q3_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q3_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q2_k_r4_ref(const float * GGML_RESTRICT x, block_q2_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_q2_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q2_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q2_k_r4(const block_q2_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q2_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q4_k_r4_ref(const float * GGML_RESTRICT x, block_q4_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_q4_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q4_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q4_k_r4(const block_q4_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q4_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q5_k_r4_ref(const float * GGML_RESTRICT x, block_q5_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_q5_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q5_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q5_k_r4(const block_q5_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q5_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q6_k_r4_ref(const float * GGML_RESTRICT x, block_q6_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_q6_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q6_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q6_k_r4(const block_q6_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q6_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq5_k_r4_ref(const float * GGML_RESTRICT x, block_iq5_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq5_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq5_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq5_k_r4(const block_iq5_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq5_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq4_k_r4_ref(const float * GGML_RESTRICT x, block_iq4_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq4_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq4_k_r4(const block_iq4_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq4_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq3_k_r4_ref(const float * GGML_RESTRICT x, block_iq3_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq3_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq3_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq3_k_r4(const block_iq3_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq3_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_k_r4_ref(const float * GGML_RESTRICT x, block_iq2_k_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_k_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_k_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq2_k_r4(const block_iq2_k_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq2_k_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq4_ks_r4_ref(const float * GGML_RESTRICT x, block_iq4_ks_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq4_ks_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq4_ks_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq4_ks_r4(const block_iq4_ks_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq4_ks_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq5_ks_r4_ref(const float * GGML_RESTRICT x, block_iq5_ks_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq5_ks_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq5_ks_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq5_ks_r4(const block_iq5_ks_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq5_ks_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_xxs_r4_ref(const float * GGML_RESTRICT x, block_iq2_xxs_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_xxs_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_xxs_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq2_xxs_r4(const block_iq2_xxs_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq2_xxs_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_xs_r4_ref(const float * GGML_RESTRICT x, block_iq2_xs_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_xs_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_xs_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq2_xs_r4(const block_iq2_xs_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq2_xs_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq2_s_r4_ref(const float * GGML_RESTRICT x, block_iq2_s_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq2_s_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq2_s_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq2_s_r4(const block_iq2_s_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq2_s_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq3_xxs_r4_ref(const float * GGML_RESTRICT x, block_iq3_xxs_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq3_xxs_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq3_xxs_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq3_xxs_r4(const block_iq3_xxs_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq3_xxs_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq3_s_r4_ref(const float * GGML_RESTRICT x, block_iq3_s_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq3_s_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq3_s_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq3_s_r4(const block_iq3_s_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq3_s_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq1_s_r4_ref(const float * GGML_RESTRICT x, block_iq1_s_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq1_s_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq1_s_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq1_s_r4(const block_iq1_s_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq1_s_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_iq1_m_r4_ref(const float * GGML_RESTRICT x, block_iq1_m_r4 * GGML_RESTRICT y, int64_t k); +void quantize_row_iq1_m_r4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_iq1_m_r4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_iq1_m_r4(const block_iq1_m_r4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_iq1_m_r4_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q8_k_r8_ref(const float * GGML_RESTRICT x, block_q8_k_r8 * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_k_r8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q8_k_r8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q8_k_r8(const block_q8_k_r8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q8_k_r8_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q8_k_r16_ref(const float * GGML_RESTRICT x, block_q8_k_r16 * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_k_r16(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q8_k_r16(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q8_k_r16(const block_q8_k_r16 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q8_k_r16_q8_k(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q8_KV_ref(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_KV(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q8_KV(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q8_KV(const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q8_KV_q8_KV(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q8_KV_r8_ref(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_KV_r8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q8_KV_r8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q8_KV_r8(const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q8_KV_r8_q8_KV(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void quantize_row_q1_0_g128_ref(const float * GGML_RESTRICT x, block_q1_0_g128 * GGML_RESTRICT y, int64_t k); +void quantize_row_q1_0_g128(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +size_t quantize_q1_0_g128(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * use_data); +void dequantize_row_q1_0_g128(const block_q1_0_g128 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +void vec_dot_q1_0_g128_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +void iqk_quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); +void quantize_row_q8_K64_ref(const float * GGML_RESTRICT x, block_q8_K64 * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_K64(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_K128(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_K16(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_K32(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_KR8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_0_x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_1_x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q8_2_x4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void iqk_quantize_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + +void repack_f32_bf16_r16 (const void * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row); +void repack_bf16_bf16_r16(const void * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row); + +void iqk_repack_tensor(struct ggml_tensor * tensor); +bool iqk_modify_tensor(struct ggml_tensor * tensor); + +int iqk_repacked_type(const struct ggml_tensor * tensor); // int instead of ggml_type so we don't need to include ggml.h +bool iqk_should_modify_tensor(const struct ggml_tensor * tensor); + +// So we can re-pack Microsoft's BitNet I2_S quants +void dequantize_row_ms_i2s(const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +typedef void (*to_float_t) (const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +typedef void (*from_float_t)(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void iqk_quantize_any(int from_type, int to_type, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, + uint64_t nb0, uint64_t nb1, uint64_t nb2, uint64_t nb3, + const void * GGML_RESTRICT x, void * GGML_RESTRICT y, void * work_buffer, + to_float_t to_float, from_float_t from_float, int ith, int nth); + +bool iqk_validate_tensor(const struct ggml_tensor * src); + +// opencoti F5-opt W2 (#290): ik_llama's base quantizers take a trailing 6th arg +// (const quantize_user_data *) that llamafile's ggml-quants.h does NOT have on the +// same names — declaring the 6-arg form on the bare name would be a "conflicting +// declaration" against the 5-arg upstream prototypes. So the iqk side renames these +// calls to iqkbase_ (sed in iqk_quantize.cpp) and we declare the 6-arg +// variants here. (iqk_ was avoided because iqk_quantize_q4_0 already names a +// pre-existing 3-arg row helper.) quantize_q6_0 / iq1{s,m}_process_1block are IK-only +// (absent upstream). Resolving these symbols is a later BUILD.mk/dispatch concern; +// this only needs to compile. +size_t iqkbase_quantize_q4_0 (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_q5_0 (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_q6_0 (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_q8_0 (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_q2_K (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_q3_K (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_q4_K (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_q5_K (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_q6_K (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_iq2_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_iq2_xs (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_iq2_s (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_iq3_xxs(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_iq3_s (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_iq4_nl (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +size_t iqkbase_quantize_iq4_xs (const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix, const struct quantize_user_data * user_data); +void iqkbase_iq1s_process_1block(int block_size, const float * xb, const float * weight, int8_t * L, float * the_scale, uint16_t * the_index, int * the_shift, float * pairs, float * sumx, float * sumw); +void iqkbase_iq1m_process_1block(const float * xb, const float * weight, int8_t * L, float * the_scale, uint16_t * the_index, int * the_shift, float * pairs); + +#ifdef __cplusplus +} +#endif diff --git a/llama.cpp/ggml/src/iqk/iqk_utils.h b/llama.cpp/ggml/src/iqk/iqk_utils.h new file mode 100644 --- /dev/null +++ b/llama.cpp/ggml/src/iqk/iqk_utils.h @@ -0,0 +1,246 @@ +#pragma once + +#include "iqk_config.h" + +#if defined IQK_IMPLEMENT + +#include "ggml-impl.h" + +#if defined(__ARM_NEON) && defined(__aarch64__) +// copy-pasted from Justine Tunney's contribution to llama.cpp +// adapted from arm limited optimized routine +// the maximum error is 1.45358 plus 0.5 ulps +// numbers above 88.38 will flush to infinity +// numbers beneath -103.97 will flush to zero +static inline float32x4_t v_expf(float32x4_t x) { + const float32x4_t r = vdupq_n_f32(0x1.8p23f); + const float32x4_t z = vfmaq_f32(r, x, vdupq_n_f32(0x1.715476p+0f)); + const float32x4_t n = vsubq_f32(z, r); + const float32x4_t b = vfmsq_f32(vfmsq_f32(x, n, vdupq_n_f32(0x1.62e4p-1f)), n, + vdupq_n_f32(0x1.7f7d1cp-20f)); + const uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_f32(z), 23); + const float32x4_t k = vreinterpretq_f32_u32(vaddq_u32(e, vreinterpretq_u32_f32(vdupq_n_f32(1)))); + const uint32x4_t c = vcagtq_f32(n, vdupq_n_f32(126)); + const float32x4_t u = vmulq_f32(b, b); + const float32x4_t j = vfmaq_f32( + vmulq_f32(vdupq_n_f32(0x1.ffffecp-1f), b), + vfmaq_f32(vfmaq_f32(vdupq_n_f32(0x1.fffdb6p-2f), vdupq_n_f32(0x1.555e66p-3f), b), + vfmaq_f32(vdupq_n_f32(0x1.573e2ep-5f), vdupq_n_f32(0x1.0e4020p-7f), b), u), u); + if (!vpaddd_u64(vreinterpretq_u64_u32(c))) + return vfmaq_f32(k, j, k); + const uint32x4_t d = vandq_u32(vclezq_f32(n), vdupq_n_u32(0x82000000)); + const float32x4_t s1 = vreinterpretq_f32_u32(vaddq_u32(d, vdupq_n_u32(0x7f000000))); + const float32x4_t s2 = vreinterpretq_f32_u32(vsubq_u32(e, d)); + return vbslq_f32(vcagtq_f32(n, vdupq_n_f32(192)), vmulq_f32(s1, s1), + vbslq_f32(c, vmulq_f32(vfmaq_f32(s2, s2, j), s1), vfmaq_f32(k, k, j))); +} +static inline float16x8_t v_expf(float16x8_t x) { + auto val1 = v_expf(vcvt_f32_f16(vget_low_f16(x))); + auto val2 = v_expf(vcvt_f32_f16(vget_high_f16(x))); + return vcombine_f16(vcvt_f16_f32(val1), vcvt_f16_f32(val2)); +} +static inline float32x4_t v_tanh(float32x4_t x) { + const float32x4_t one = vdupq_n_f32(1.0f); + const float32x4_t two_x = vmulq_f32(x, vdupq_n_f32(2.f)); + const float32x4_t exp_two_x = v_expf(two_x); + const uint32x4_t mask = vcgtq_f32(x, vdupq_n_f32(10.f)); + const float32x4_t res = vdivq_f32(vsubq_f32(exp_two_x, one), vaddq_f32(exp_two_x, one)); + return vreinterpretq_f32_u32(vorrq_u32(vandq_u32(vreinterpretq_u32_f32(one), mask), vbicq_u32(vreinterpretq_u32_f32(res), mask))); + //return vdivq_f32(vsubq_f32(exp_two_x, one), vaddq_f32(exp_two_x, one)); +} +//inline float32x4_t v_tanh(float16x8_t x) { +// auto val1 = v_tanh(vcvt_f32_f16(vget_low_f16(x))); +// auto val2 = v_tanh(vcvt_f32_f16(vget_high_f16(x))); +// return vcombine_f16(vcvt_f16_f32(val1), vcvt_f16_f32(val2)); +//} +static inline float32x4_t v_silu(float32x4_t x) { + const float32x4_t one = vdupq_n_f32(1.0f); + const float32x4_t zero = vdupq_n_f32(0.0f); + const float32x4_t neg_x = vsubq_f32(zero, x); + const float32x4_t exp_neg_x = v_expf(neg_x); + const float32x4_t one_plus_exp_neg_x = vaddq_f32(one, exp_neg_x); + return vdivq_f32(x, one_plus_exp_neg_x); +} +static inline float32x4_t v_silu_oai(float32x4_t x, float32x4_t alpha) { + const float32x4_t one = vdupq_n_f32(1.0f); + const float32x4_t neg_x = vmulq_f32(alpha, x); + const float32x4_t exp_neg_x = v_expf(neg_x); + const float32x4_t one_plus_exp_neg_x = vaddq_f32(one, exp_neg_x); + return vdivq_f32(x, one_plus_exp_neg_x); +} +static inline float32x4_t v_gelu(float32x4_t x, float32x4_t c1, float32x4_t c2) { + const float32x4_t one = vdupq_n_f32(1.0f); + float32x4_t arg = vfmaq_f32(one, c1, vmulq_f32(x, x)); + arg = vmulq_f32(arg, vmulq_f32(x, c2)); + float32x4_t exp_arg = v_expf(arg); + float32x4_t gelu = vmulq_f32(x, vdivq_f32(exp_arg, vaddq_f32(exp_arg, one))); + uint32x4_t mask = vcgtq_f32(x, vdupq_n_f32(10.f)); + return vbslq_f32(mask, x, gelu); +} + +#endif // __ARN_NEON + +#if defined(__AVX512F__) && defined(_MSC_VER) +#include + +#ifndef __clang__ +static inline __m512i operator|(__m512i a, __m512i b) { return _mm512_or_si512(a, b); } +static inline __m512i operator&(__m512i a, __m512i b) { return _mm512_and_si512(a, b); } +static inline __m512i operator^(__m512i a, __m512i b) { return _mm512_xor_si512(a, b); } +#endif +#endif + +#if defined(__AVX512F__) && defined(__AVX512DQ__) + +// copy-pasted from Justine Tunney's contribution to llama.cpp +// adapted from arm limited optimized routine +// the maximum error is 1.45358 plus 0.5 ulps +// numbers above 88.38 will flush to infinity +// numbers beneath -103.97 will flush to zero +static inline __m512 v_expf(__m512 x) { + const __m512 r = _mm512_set1_ps(0x1.8p23f); + const __m512 z = _mm512_fmadd_ps(x, _mm512_set1_ps(0x1.715476p+0f), r); + const __m512 n = _mm512_sub_ps(z, r); + const __m512 b = + _mm512_fnmadd_ps(n, _mm512_set1_ps(0x1.7f7d1cp-20f), + _mm512_fnmadd_ps(n, _mm512_set1_ps(0x1.62e4p-1f), x)); + const __mmask16 d = + _mm512_cmp_ps_mask(_mm512_abs_ps(n), _mm512_set1_ps(192), _CMP_GT_OQ); + const __m512 u = _mm512_mul_ps(b, b); + const __m512 j = _mm512_fmadd_ps( + _mm512_fmadd_ps(_mm512_fmadd_ps(_mm512_set1_ps(0x1.0e4020p-7f), b, + _mm512_set1_ps(0x1.573e2ep-5f)), + u, + _mm512_fmadd_ps(_mm512_set1_ps(0x1.555e66p-3f), b, + _mm512_set1_ps(0x1.fffdb6p-2f))), + u, + _mm512_fmadd_ps(_mm512_set1_ps(0x1.ffffecp-1f), b, _mm512_set1_ps(1.0F))); + const __m512 res = _mm512_scalef_ps(j, n); + if (_mm512_kortestz(d, d)) + return res; + const __m512 zero = _mm512_setzero_ps(); + const __m512 alt = _mm512_mask_blend_ps( + _mm512_cmp_ps_mask(n, zero, _CMP_LE_OQ), _mm512_set1_ps(INFINITY), zero); + return _mm512_mask_blend_ps(d, res, alt); +} +static inline __m512 v_tanh(__m512 x) { + const __m512 one = _mm512_set1_ps(1.0f); + const __m512 exp_two_x = v_expf(_mm512_mul_ps(x, _mm512_set1_ps(2.f))); + const __mmask16 mask = _mm512_cmp_ps_mask(x, _mm512_set1_ps(10.f), _CMP_GT_OQ); + const __m512 res = _mm512_div_ps(_mm512_sub_ps(exp_two_x, one), _mm512_add_ps(exp_two_x, one)); + return _mm512_mask_blend_ps(mask, res, one); +} +static inline __m512 v_gelu(__m512 x, __m512 c1, __m512 c2) { + const __m512 one = _mm512_set1_ps(1.0f); + __m512 arg = _mm512_fmadd_ps(x, _mm512_mul_ps(c1, x), one); + //__m512 arg = _mm512_add_ps(one, _mm512_mul_ps(_mm512_mul_ps(x, x), c1)); + arg = _mm512_mul_ps(arg, _mm512_mul_ps(c2, x)); + const __mmask16 mask = _mm512_cmp_ps_mask(arg, _mm512_set1_ps(30.f), _CMP_GT_OQ); + const __m512 exp_arg = v_expf(arg); + const __m512 ratio = _mm512_div_ps(exp_arg, _mm512_add_ps(exp_arg, one)); + return _mm512_mul_ps(x, _mm512_mask_blend_ps(mask, ratio, one)); +} +static inline __m512 v_silu(__m512 x) { + const __m512 one = _mm512_set1_ps(1); + const __m512 zero = _mm512_setzero_ps(); + const __m512 neg_x = _mm512_sub_ps(zero, x); + const __m512 exp_neg_x = v_expf(neg_x); + const __m512 one_plus_exp_neg_x = _mm512_add_ps(one, exp_neg_x); + return _mm512_div_ps(x, one_plus_exp_neg_x); +} +static inline __m512 v_silu_oai(__m512 x, __m512 alpha) { + const __m512 one = _mm512_set1_ps(1); + const __m512 neg_x = _mm512_mul_ps(alpha, x); + const __m512 exp_neg_x = v_expf(neg_x); + const __m512 one_plus_exp_neg_x = _mm512_add_ps(one, exp_neg_x); + return _mm512_div_ps(x, one_plus_exp_neg_x); +} +static inline __m512 v_clamp_max(__m512 x, __m512 max) { + auto mask = _mm512_cmp_ps_mask(x, max, _CMP_GT_OQ); + return _mm512_mask_blend_ps(mask, x, max); +} +#endif // __AVX512__ + +#if defined(__AVX2__) && defined(__FMA__) + +// adapted from arm limited optimized routine +// the maximum error is 1.45358 plus 0.5 ulps +// numbers above 88.38 will flush to infinity +// numbers beneath -103.97 will flush to zero +static inline __m256 v_expf(__m256 x) { + const __m256 r = _mm256_set1_ps(0x1.8p23f); + const __m256 z = _mm256_fmadd_ps(x, _mm256_set1_ps(0x1.715476p+0f), r); + const __m256 n = _mm256_sub_ps(z, r); + const __m256 b = _mm256_fnmadd_ps(n, _mm256_set1_ps(0x1.7f7d1cp-20f), + _mm256_fnmadd_ps(n, _mm256_set1_ps(0x1.62e4p-1f), x)); + const __m256i e = _mm256_slli_epi32(_mm256_castps_si256(z), 23); + const __m256 k = _mm256_castsi256_ps( + _mm256_add_epi32(e, _mm256_castps_si256(_mm256_set1_ps(1)))); + const __m256i c = _mm256_castps_si256( + _mm256_cmp_ps(_mm256_andnot_ps(_mm256_set1_ps(-0.f), n), + _mm256_set1_ps(126), _CMP_GT_OQ)); + const __m256 u = _mm256_mul_ps(b, b); + const __m256 j = _mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_fmadd_ps(_mm256_set1_ps(0x1.0e4020p-7f), b, + _mm256_set1_ps(0x1.573e2ep-5f)), u, + _mm256_fmadd_ps(_mm256_set1_ps(0x1.555e66p-3f), b, + _mm256_set1_ps(0x1.fffdb6p-2f))), + u, _mm256_mul_ps(_mm256_set1_ps(0x1.ffffecp-1f), b)); + if (!_mm256_movemask_ps(_mm256_castsi256_ps(c))) + return _mm256_fmadd_ps(j, k, k); + const __m256i g = _mm256_and_si256( + _mm256_castps_si256(_mm256_cmp_ps(n, _mm256_setzero_ps(), _CMP_LE_OQ)), + _mm256_set1_epi32(0x82000000u)); + const __m256 s1 = + _mm256_castsi256_ps(_mm256_add_epi32(g, _mm256_set1_epi32(0x7f000000u))); + const __m256 s2 = _mm256_castsi256_ps(_mm256_sub_epi32(e, g)); + const __m256i d = _mm256_castps_si256( + _mm256_cmp_ps(_mm256_andnot_ps(_mm256_set1_ps(-0.f), n), + _mm256_set1_ps(192), _CMP_GT_OQ)); + return _mm256_or_ps( + _mm256_and_ps(_mm256_castsi256_ps(d), _mm256_mul_ps(s1, s1)), + _mm256_andnot_ps( + _mm256_castsi256_ps(d), + _mm256_or_ps( + _mm256_and_ps(_mm256_castsi256_ps(c), + _mm256_mul_ps(_mm256_fmadd_ps(s2, j, s2), s1)), + _mm256_andnot_ps(_mm256_castsi256_ps(c), _mm256_fmadd_ps(k, j, k))))); +} +static inline __m256 v_tanh(__m256 x) { + const __m256 one = _mm256_set1_ps(1.0f); + const __m256 exp_two_x = v_expf(_mm256_mul_ps(x, _mm256_set1_ps(2.f))); + const __m256 res = _mm256_div_ps(_mm256_sub_ps(exp_two_x, one), _mm256_add_ps(exp_two_x, one)); + const __m256 mask = _mm256_cmp_ps(x, _mm256_set1_ps(10.f), _CMP_GT_OQ); + return _mm256_or_ps(_mm256_and_ps(mask, one), _mm256_andnot_ps(mask, res)); +} +static inline __m256 v_gelu(__m256 x, __m256 c1, __m256 c2) { + const __m256 one = _mm256_set1_ps(1.0f); + const __m256 mask = _mm256_cmp_ps(x, _mm256_set1_ps(10.f), _CMP_GT_OQ); + __m256 arg = _mm256_add_ps(one, _mm256_mul_ps(_mm256_mul_ps(x, x), c1)); + arg = _mm256_mul_ps(arg, _mm256_mul_ps(x, c2)); + __m256 exp_arg = v_expf(arg); + __m256 gelu = _mm256_mul_ps(x, _mm256_div_ps(exp_arg, _mm256_add_ps(exp_arg, one))); + return _mm256_or_ps(_mm256_and_ps(mask, x), _mm256_andnot_ps(mask, gelu)); +} +static inline __m256 v_silu(__m256 x) { + const __m256 one = _mm256_set1_ps(1); + const __m256 zero = _mm256_setzero_ps(); + const __m256 neg_x = _mm256_sub_ps(zero, x); + const __m256 exp_neg_x = v_expf(neg_x); + const __m256 one_plus_exp_neg_x = _mm256_add_ps(one, exp_neg_x); + return _mm256_div_ps(x, one_plus_exp_neg_x); +} +static inline __m256 v_silu_oai(__m256 x, __m256 alpha) { + const __m256 one = _mm256_set1_ps(1); + const __m256 neg_x = _mm256_mul_ps(alpha, x); + const __m256 exp_neg_x = v_expf(neg_x); + const __m256 one_plus_exp_neg_x = _mm256_add_ps(one, exp_neg_x); + return _mm256_div_ps(x, one_plus_exp_neg_x); +} +static inline __m256 v_clamp_max(__m256 x, __m256 max) { + auto mask = _mm256_cmp_ps(x, max, _CMP_GT_OQ); + return _mm256_or_ps(_mm256_and_ps(mask, max), _mm256_andnot_ps(mask, x)); +} + +#endif // __AVX2__ + +#endif // IQK_IMPLEMENT